PytorchLightningフレームワークでDDPを実装
PytorchLightning 1.4.0まではTrainer(accelerator="ddp")
PytorchLightning 1.5.0からはTrainer(strategy="ddp")推奨に変更された。
このように破壊的変更がまだまだ存在することに注意
さらにv1.5以降、strategy="ddp"、device="autoにすることで空いているGPUを用いてDDPを行うことが容易になった
やり方
strategy="gpu"を設定する
例：trainer = Trainer(gpus=8, strategy="ddp")
これによってスクリプトを内部で複数回呼び出す形式になる
とても簡単
DPやDDP2分散モードを使用する場合
DDPと異なり、各々のGPUに渡りバッチを分割する場合
DPはDDPとくらべてメインGPUに各トレーニング結果を集約する必要がある。
各トレーニング結果を手動で集めるために
xxx_step_end()ですべてのデバイスから出力を取得して集約するコードを書く必要がある。
ここで初めて、training_step()とtraining_step_end()とを分けた意味がある。ふつうはstep_end()は使わずにstep()の最後に書けばよい。
ステップ・が終わるごとにlossなどのlog結果を集約する必要がある。
lossを平均化する場合は以下のように書く必要がある。
code:lit_train.py
def training_step(self, batch, batch_idx):
x, y = batch
y_hat = self.model(x)
loss = F.cross_entropy(y_hat, y)
pred = ...
return {"loss": loss, "pred": pred}
def training_step_end(self, batch_parts):
# predictions from each GPU
predictions = batch_parts"pred" # losses from each GPU
gpu_0_prediction = predictions0 gpu_1_prediction = predictions1 # do something with both outputs
return (losses0 + losses1) / 2 def training_epoch_end(self, training_step_outputs):
for out in training_step_outputs:
...