勾配蓄積
大きなバッチサイズを限られたGPUメモリでやりくりするために勾配を貯める。
実質上のバッチサイズを増やす機構。
「RuntimeError: CUDA out of memory 等が出ているけれど、バッチサイズは増やしたい。」
そんなわがままを叶えることができる。
ただしGPUのI/Oが増えるから訓練時間も増える。
PyTorchでなぜ勾配を0にする処理(optimizer.zero_grad())を手動で設定するかといえば、勾配蓄積を行えるようにするためである。
こんな感じで書く
code:train.py
for i, (inputs, labels) in tqdm(enumerate(data_loader), total=len(data_loader), ncols=80, desc="train"):
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
# Gradient accumulation
# grad_acc (int): 実質上バッチサイズを何倍に増やすか. デフォルトは1.
# 大きなバッチサイズを限られたGPUメモリでやりくりするために勾配を貯める
if (i % grad_acc) == 0:
optimizer.step() # オプティマイザによる学習パラメータの更新
optimizer.zero_grad() # 勾配を0にする
ほかにもgradient accumulation pytorchとかで検索すると良さげ資料が出るかも