【PyTorch】Mixed Precision (autocast) での an illegal memory access was encountered / 処理が遅い問題

PyTorchのバージョンは1.13.0です．

PyTorchのMixed Precisionを試した時に起こった現象と解決した方法についてのメモです． 現在のPyTorchでMixed Precisionを行うには，以下のように with torch.autocast("cuda", torch.float16): で囲むことで勝手に半精度で計算してくれます．

pytorch.org

エラー

ランダムな処理を行うプログラムを回していたところ，以下の2種類のエラーが出たり出なかったりしました．

RuntimeError: CUDA error: CUBLAS_STATUS_EXECUTION_FAILED when calling `cublasGemmStridedBatchedExFix

RuntimeError: CUDA error: an illegal memory access was encountered terminate called after throwing an instance of 'c10::Error' what(): CUDA error: an illegal memory access was encountered

このエラーが出た個所がF.affine_gridを呼び出しているところだったためそれ関連のエラーかと思いましたが，最終的には以下のissueに答えがありました．

github.com

どうやら単純にVRAMが足りなかっただけだったようです．実際，バッチサイズを減らしてみると見事にエラーが出なくなりました．かなしい

(この直前にランダムに画像をパディングする処理をかけており，パディング幅が大きい時にF.affine_gridを呼び出した時点でVRAM不足になったのだと思います)

処理が遅い問題

次に起こったのは処理が一向に進まないという問題でした． この解決方法は，「読み込んだデータを.to("cuda")でキャストする」でした．

↓のような感じです

...
with torch.autocast("cuda", torch.float16):
    x = x.to("cuda")   # この行を追加
    y = model(x)
    ...

上記のwith文の中ではデータを自動でキャストしてくれるはずなので.to("cuda")は無くてもエラーは起こらないのですが，なぜか処理が進みませんでした．

(この辺りについて詳しい方がいればコメント頂けると幸いです🙇‍♂️)