RTX5070を買ったので、簡単にできるPyTorchの高速化手法を試してみた

新しくGeForce RTX 5070を買ったので、PyTorch公式チュートリアルCIFAR10データセットの画像分類でベンチを取りました。

結論から言うとモデルのトレーニングの開始から完了までで、73.17秒かかりました。

まあ、そこそこ早いですね。もっと早くできないか気になったため、高速化を試みます。

# 高速化された手法
batchsizeを2から128に変更 3.49!!!!
pin_memory=Trueを設定 52.18

# 逆に遅くなった手法
numworkersを2から8に変更 76.02
torch.backends.cudnn.benchmark=True　一回目 75.75
torch.backends.cudnn.benchmark=True　二回目 79.51
AMP　87.23

batchsizeを大きくするとなんと20倍以上高速化しました。
pin_memory=Trueを設定するのも効果がありそうです。
それ以外の高速化手法は逆に学習時間が伸びています。

データセットが小さい上に、epochも２なので実践的なベンチではありませんが、これだけ学習速度が変わってくるなら、日ごろからbatchsizeは適切な値にするようおに心がけたほうがいいですね。

まとめ

PyTorchの高速化は、何よりも適切なbatchsizeが重要。
pin_memory=Trueも効果あり。