🔰

SageMaker の推論処理を理解する：predict_fn の役割と使い方（PyTorch）

に公開
AWS
PyTorch
初心者向け
SageMaker
学びの記録
tech

はじめに

これまで、SageMaker における model_fn と input_fn の役割について整理してきました。
今回はその続きとして、モデルを用いた推論処理を担う predict_fn に焦点を当ててみます。

個人の学習記録として、現時点での理解や気づきをまとめておきます。

predict_fn とは？

predict_fn は、前処理された入力データとモデルを受け取り、予測（推論）を実行する関数です。

def predict_fn(input_data, model):
    return model(input_data)

PyTorch の場合、model(input_data) のようにシンプルに記述することも可能ですが、
実運用では torch.no_grad() の使用が推奨されるケースが多くなります。

典型的な処理の流れ（PyTorchの場合）

def predict_fn(input_data, model):
    try:
        # 入力データの基本チェック
        if input_data is None:
            raise ValueError("input_data is None")
        
        # モデルの基本チェック
        if model is None:
            raise ValueError("model is None")
        
        with torch.no_grad():  # 勾配を計算しない（高速化・メモリ削減）
            output = model(input_data)
        
        return output
    
    except Exception as e:
        # エラーをログに出力（SageMakerのログで確認可能）
        print(f"予測処理でエラーが発生しました: {e}")
        raise  # エラーを再発生させてSageMakerに通知

このように、predict_fn では学習時とは異なる「推論時の設定」に注意する必要があります。

どのタイミングで呼ばれる？

SageMakerでは、以下のような流れで関数が呼ばれます。

リクエスト受信 → input_fn → predict_fn → output_fn → レスポンス返却

つまり、input_fn によって整形されたデータが predict_fn に渡され、
そこでモデルを用いた実際の予測処理が実行される、という位置づけになります。

設計のポイント・注意点

  • torch.no_grad() を使って余分なメモリ消費を防ぐ
  • 推論の出力形式を意識する（Softmaxなどで確率に変換するか？など）
  • 必要に応じて追加の後処理（ラベルID → ラベル名など）もここで行うことができる

今の時点での自分の理解（仮）

  • predict_fn は「モデルに入力を渡して結果を返す」メイン処理
  • 入力とモデルは input_fn / model_fn から受け取る前提
  • モデルによっては追加の整形処理（出力の変換など）も必要になりそう

おわりに

predict_fn は、モデルを実際に動かして予測を行う中心的な役割を担う関数です。
今回改めて整理することで、SageMakerにおける「入力から出力までの流れ」が、少しずつ明確になってきたように感じています。

次回は、レスポンス形式を整える役割を担う output_fn についても整理する予定です。

🔗 関連記事

Discussion