PyTorch Lightning 1.1: notebooks : Lightning フラグへのイントロダクション (翻訳/解説)
翻訳 : (株)クラスキャット セールスインフォメーション
作成日時 : 02/17/2021 (1.1.x)

* 本ページは、PyTorch Lightning ドキュメントの以下のページの前半を翻訳した上で適宜、補足説明したものです：

• notebooks : Introduction to Lightning Flags

* サンプルコードの動作確認はしておりますが、必要な場合には適宜、追加改変しています。
* ご自由にリンクを張って頂いてかまいませんが、sales-info@classcat.com までご一報いただけると嬉しいです。

 

 

notebooks : Lightning フラグへのイントロダクション

• Trainer フラグ – 利用可能な Lightning Trainer フラグの概要です。

このノートブックでは、Trainer オブジェクトで利用可能なフラグを調べます。総てが Colab 環境で動作はしないことに注意してください (マルチ GPU, 等)。

 

セットアップ

まずは最初に大事なこととして、Lightning をインストールする必要があります。単純に pip install pytorch-lightning

! pip install pytorch-lightning

import os

from argparse import ArgumentParser
import torch
from torch import nn
from torch.nn import functional as F
from torch.utils.data import DataLoader
from torch.utils.data import random_split
from torchvision.datasets import MNIST
from torchvision import transforms
import pytorch_lightning as pl
from pytorch_lightning.metrics.functional import accuracy

from torchvision.datasets.mnist import MNIST
from torchvision import transforms

# ------------
# data
# ------------
pl.seed_everything(1234)
batch_size = 32

# Init DataLoader from MNIST Dataset

dataset = MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor())
mnist_test = MNIST(os.getcwd(), train=False, download=True, transform=transforms.ToTensor())
mnist_train, mnist_val = random_split(dataset, [55000, 5000])

train_loader = DataLoader(mnist_train, batch_size=batch_size)
val_loader = DataLoader(mnist_val, batch_size=batch_size)
test_loader = DataLoader(mnist_test, batch_size=batch_size)

 

単純なオートエンコーダ・モデル

単純な Lightning モデルを定義していき Lightning Trainer の総ての設定で遊ぶことができます。

LightningModule はフックに再構造化された単純に純粋な PyTorch で、訓練プロセスの総てのステップを表します。

貴方のモデルの総てのパートを制御するために LightningModule フックを使用できますが、このノートブックの目的のために非常に単純な MNIST 分類器、手書き画像の 28*28 グレースケール画像を取り、そして 0-9 間の数字を予測できるモデルを使用します。

LightningModule は画像分類器のような単一モデル、あるいはエンコーダとデコーダを含むこのオートエンコーダのような複数のモデルから成る深層学習システムを包含できます。

class LitAutoEncoder(pl.LightningModule):

    def __init__(self, batch_size=32, lr=1e-3):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(28 * 28, 64),
            nn.ReLU(),
            nn.Linear(64, 3)
        )
        self.decoder = nn.Sequential(
            nn.Linear(3, 64),
            nn.ReLU(),
            nn.Linear(64, 28 * 28)
        )
        self.batch_size=batch_size
        self.learning_rate=lr

    def forward(self, x):
        # in lightning, forward defines the prediction/inference actions
        embedding = self.encoder(x)
        return embedding

    def training_step(self, batch, batch_idx):
        x, y = batch
        x = x.view(x.size(0), -1)
        z = self.encoder(x)
        x_hat = self.decoder(z)
        loss = F.mse_loss(x_hat, x)
        self.log('train_loss', loss)
        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        x = x.view(x.size(0), -1)
        z = self.encoder(x)
        x_hat = self.decoder(z)
        loss = F.mse_loss(x_hat, x)
        self.log('val_loss', loss)
        
    def test_step(self, batch, batch_idx):
        x, y = batch
        x = x.view(x.size(0), -1)
        z = self.encoder(x)
        x_hat = self.decoder(z)
        loss = F.mse_loss(x_hat, x)
        self.log('test_loss', loss)

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
        return optimizer

LightningModule はエポックとバッチループを持たず、model.train() と model.eval() を呼び出しておらず、そして CUDA やハードウェアの言及を持たないことに気付くでしょう。それはそれは総て Lightning Trainer により自動化されるからです。総てのエンジニアリング・ボイラープレートは trainer により自動化されます :

• 訓練ループ
• 評価とテストループ
• 適切なときに model.train(), model.eval(), no_grad を呼び出す
• CUDA or to_device 呼び出し

それはコードの変更なしに GPU と TPU のような異なるハードウェア上で貴方のモデルを訓練することを可能にします！

 

lightning trainer を単純に使用するために

1. LightningModule とデータセットを初期化する
2. lightning trainer を初期化する
3. trainer.fit を呼び出す

#####################
# 1. Init Model
#####################

model = LitAutoEncoder()

#####################
# 2. Init Trainer
#####################

# these 2 flags are explained in the later sections...but for short explanation:
# - progress_bar_refresh_rate: limits refresh rate of tqdm progress bar so Colab doesn't freak out
# - max_epochs: only run 2 epochs instead of default of 1000
trainer = pl.Trainer(progress_bar_refresh_rate=20, max_epochs=2)

#####################
# 3. Train
#####################
trainer.fit(model, train_loader, val_loader)

Lightning デフォルトを使用して、モデルは丁度そのように訓練されています。Lightning の美点は総てが容易に configurable であることです。次のノートブックでは、訓練・検証とテストループを制御し、GPU と TPU 上で実行し、チェックポインティング、early stopping、それ以上のようなことを行なうために貴方の Trainer を制御できる総ての方法を示していきます。

 

訓練ループと eval ループフラグ

ネットワークを実際にスケールアップするために、GPU のようなアクセラレータを使用することができます。GPU あるいはグラフィカル・プロセッシング・ユニットは行列操作を並列化して、CPU 上の訓練を越えて少なくとも 100x のスピードアップを可能にします。

8 GPU を持つマシンを持つとしましょう。このフラグを 1, 4 か 8 GPU に設定することができて lightning は貴方のために訓練を自動的に分散します。

trainer = pl.Trainer(gpus=1)

 
Lightning は貴方のコードをハードウェア不可知にします… これはコード変更なしに CPU, GPU 間を切り替えることができることを意味します。

けれども、それは良い PyTorch 習慣を形成することを必要とします :

1. 最初に、貴方のコード内の .cuda() or .to() 呼び出しを除去します。
2. 2 番目に、新しい tensor を初期化するとき、呼び出しで device=self.device を設定します、何故ならば総ての lightningModule はそれが何の gpu インデックスや TPU コアの上にあるか知っているからあです。

type_as を使用することもできます、そしてあるいは register_buffer() で tensor をモジュールの init メソッド内のバッファとして登録することもできます。

# before lightning
def forward(self, x):
    z = torch.Tensor(2, 3)
    z = z.cuda(0)

# with lightning
def forward(self, x):
    z = torch.Tensor(2, 3)
    z = z.type_as(x, device=self.device)

class LitModel(LightningModule):

    def __init__(self):
        ...
        self.register_buffer("sigma", torch.eye(3))
        # you can now access self.sigma anywhere in your module

Lightning Trainer はエポックとバッチに渡る反復、 訓練・評価とテストループ、CUDA と to(device) 呼び出し、model.train と model.eval 呼び出しのような総てのエンジニアリング・ボイラープレートを自動化します。

以下の trainer フラグを使用することにより、貴方は依然としてループに渡る完全な制御を持ちます :

 

検証ステップを呼び出す

時に、エポックの訓練は非常に高速かもしれません、エポック毎数分のように。この場合、総てのエポック上で検証する必要はないかもしれません。代わりに、実際には数エポック後に検証することができます。

検証ステップの頻度を制御するには check_val_every_n_epoch フラグを使用します :

# run val loop every 10 training epochs
trainer = pl.Trainer(check_val_every_n_epoch=10)

trainer.fit(model, train_loader, val_loader)

 

val_check_interval

エポックが非常に長いような幾つかのケースでは、エポック内で検証を確認したいかもしれません。

val_check_interval フラグを設定することにより訓練エポック内で検証ステップを実行することもできます。

訓練エポック内部で検証セットを確認するためには val_check_interval を [0.0 to 1.0] 間の float に設定します。例えば、それを 0.25 に設定することは訓練エポックの間に 4 回検証セットを確認します。

デフォルトは 1.0 に設定されています。

# check validation set 4 times during a training epoch
trainer = pl.Trainer(val_check_interval=0.25)

trainer.fit(model, train_loader, val_loader)

iteratable データセットを持つとき、あるいはプロダクション・ユースケースのためにデータをストリームするとき、総てのステップ数毎に検証セットを確認することは有用です。val_check_interval を int に設定します :

# check validation set every 1000 training batches
# use this when using iterableDataset and your dataset has no length
# (ie: production cases with streaming data)
trainer = pl.Trainer(val_check_interval=1000)

trainer.fit(model, train_loader, val_loader)

 

num_sanity_val_steps

検証ループにバグを持つところでは課題にぶつかるかもしれませんが、訓練ループが終了するまでそれを捕捉できません。

そして訓練ループが数時間か数日かかる場合、価値ある計算を浪費します。

代わりに、lightning はこれらの種類のバグを前もって捕捉するために最初に 2 ステップの検証を自動的に実行します。

num_sanity_val_steps フラグは訓練ルーチンを開始する前に検証の n バッチを実行するのを手助けできます。

それを無効にするにはそれを 0 に設定できます。

# turn it off
trainer = pl.Trainer(num_sanity_val_steps=0)

trainer.fit(model, train_loader, val_loader)

訓練前に総ての検証データを確認するにはそれを -1 に設定します。

# check all validation data
trainer = pl.Trainer(num_sanity_val_steps=-1)

trainer.fit(model, train_loader, val_loader)

あるいは任意の数の検証ステップを使用します。

trainer = pl.Trainer(num_sanity_val_steps=10)

trainer.fit(model, train_loader, val_loader)

 

訓練、検証とテストバッチを制限する

どのくらいの訓練、検証とテストデータセットをモデルが確認することを望むかの制限を設定できます。実際に非常に大規模な検証やテストセットを持つ場合にこれはエポックの最後に発生するデバッグやテストのようなもののために有用です。実行するバッチ数を指定するためにこのフラグを int に設定します。

# run for only 10 batches
trainer = pl.Trainer(limit_test_batches=10)

trainer.fit(model, train_loader, val_loader)

例えば、AUC ROC のような幾つかのメトリクスは検証結果全体の上で計算される必要があります。

trainer = pl.Trainer(limit_val_batches=10)

trainer.fit(model, train_loader, val_loader)

バッチを総てのエポックでセットのパーセンテージに制限するために float を使用できます。

# run through only 25% of the test set each epoch
trainer = pl.Trainer(limit_test_batches=0.25)

trainer.fit(model, train_loader, val_loader)

 

GPU 上で訓練する

1 GPU 上で実行するにはフラグを 1 に設定します。

trainer = pl.Trainer(gpus=1)

trainer.fit(model, train_loader, val_loader)

2 or 4 GPU 上で実行するには、フラグを 2 か 4 に設定します。

trainer = pl.Trainer(gpus=2)

trainer.fit(model, train_loader, val_loader)

[1, 4] のようなインデックスのリストを使用して、どの GPU デバイス上で実行するか選択することもできます。

あるいは ‘1, 2’ のような GPU id のカンマ区切りリストを含む文字列。

# list: train on GPUs 1, 4 (by bus ordering)
# trainer = Trainer(gpus='1, 4') # equivalent
trainer = pl.Trainer(gpus=[1, 4])

trainer.fit(model, train_loader, val_loader)

trainer = pl.Trainer(gpus=list(range(4)))

trainer.fit(model, train_loader, val_loader)

gpus=-1 を設定することにより利用可能な総ての GPU を利用できます。

# trainer = Trainer(gpus='-1') - equivalent
trainer = pl.Trainer(gpus=-1)

trainer.fit(model, train_loader, val_loader)

Lightning は GPU を順序付けるために PCI bus_id をインデックスとして使用します。

 

auto_select_gpus

“exclusive (排他的) モード” で実行することにより GPU を節約できます、これは同時に 1 つのプロセスだけがそれらにアクセスできることを意味します。exclusive モードで実行するときどの GPU を使用するべきか確かでない場合、Lightning は空いている GPU を貴方のために自動的に見つけることができます。

単純に gpu の数を整数 gpus=k として指定し、そして trainer フラグ auto_select_gpus=True を設定します。Lightning は他のプロセスにより専有されていない k gpu を自動的に見つける手助けをします。

# enable auto selection (will find two available gpus on system)
trainer = pl.Trainer(gpus=2, auto_select_gpus=True)

trainer.fit(model, train_loader, val_loader)

 

GPU 使用方法を分析する

log_gpu_memory

これは GPU のメモリ使用量を分析するために有用です。

マスター・ノード上の総ての GPU についての GPU メモリ使用量を得るには、フラグを log_gpu_memory=all に設定します。

内部的には、lightning は nvidia-smi コマンドを使用します、これは訓練をスローダウンさせるかもしれません。

多くの GPU からの使用料を一度にログ記録する場合ログは過剰になる可能性があります。この場合、フラグを min_max に設定することもできます、これはマスターノードの総ての GPU に渡る最小と最大使用量だけをログ記録します。

lightning はパフォーマンス的な理由で総てのノードに渡り使用量をログ記録はしないことに注意してください。

# log all the GPUs (on master node only)
trainer = Trainer(log_gpu_memory='all')

trainer.fit(model, train_loader, val_loader)

パフォーマンスの低下を避けるためにマスターノード上の最初と最大メモリだけをログ記録するために log_gpu_memory=min_max を設定することもできます。

# log only the min and max memory on the master node
trainer = Trainer(log_gpu_memory='min_max')

trainer.fit(model, train_loader, val_loader)

しかし単に一つではなくマルチマシン上で訓練することを望む場合にはどうでしょう？

 

マルチ GPU 上で訓練する

Lightning はモデルをハードウェア不可知にし、そしてフラグの小さい変更で GPU 上で実行できます。Lightning はまた多くのマシンに渡るマルチ GPU 上の訓練もサポートします。

num_nodes フラグを設定することでこれを行なうことができます。

world サイズ、あるいは使用している GPU の総数は gpus*num_nodes になります。

gpus=8 と num_nodes=32 を設定すれば 256 GPU 上で訓練していきます。

trainer = pl.Trainer(gpus=8, num_nodes=32)

trainer.fit(model, train_loader, val_loader)

 

アクセラレータ

内部的には、Lightning は GPU に渡る訓練を分散させるために分散データ並列 (DDP, distributed data parallel) を使用します。

DDP の Lightning 実装は内部的には貴方のスクリプトを正しい環境変数で複数回呼び出します。

内部的にはそれはあたかもスクリプトを以下のように呼び出したかのようにです :

1. 各ノードに渡る各 GPU はそれ自身のプロセスを得ます。
2. 各 GPU はデータセット全体のサブセットへの可視性を得ます。それは絶対にそのサブセットだけを見ます。
3. 各プロセスはモデルを初期化します。(各モデルが同じ重みで初期化するようにランダムシードを設定することを確実にします。)
4. 各プロセスは完全な forward と backward パスを並列に遂行します。
5. 勾配は総てのプロセスに渡り同期されて平均されます。
6. 各プロセスはその optimizer を更新します。モードを設定することなくマルチ GPU やノードを要求する場合、DDP が自動的に利用されます。

# ddp = DistributedDataParallel
# trainer = pl.Trainer(gpus=2, num_nodes=2) equivalent
trainer = pl.Trainer(gpus=2, num_nodes=2, accelerator='ddp')

trainer.fit(model, train_loader, val_loader)

DDP は訓練を分散させる最速で推奨される方法ですが、DDP がサポートされないとき、accelerator trainer フラグに他のバックエンドを渡すことができます。

DDP は以下では利用可能ではありません :

• Jupyter Notebook, Google COLAB, Kaggle, etc.
• root パッケージなしにネストされたスクリプトを持つ場合
• あるいはスクリプトが .fit や .test を複数回起動する必要がある場合

 

DDP_SPAWN

これらのケースでは、代わりに ddp_spawn を利用できます。ddp_spawn はそれが訓練プロセスを開始するために .spawn() を使用することを除いて正確に DDP のようなものです。

trainer = pl.Trainer(gpus=2, num_nodes=2, accelerator='ddp_spawn')

trainer.fit(model, train_loader, val_loader)

この使用を「強く」反対します、何故ならばそれは (Python と PyTorch により) 制限を持つからです :

• .spawn() はモデルをサブプロセスで訓練しますので、メインプロセス上のモデルは更新されません。
• Dataloader(num_workers=N)、ここで N は大きいです、は DDP による訓練のボトルネックになります … ie: それは「非常に」遅いか全く動作しないでしょう。これは PyTorch 制限です。
• 総てを picklable として強制します。

DDP は DDP_spawn よりも遥かに高速です。 DDP を利用できるために以下を勧めます :

1. setup.py を使用した貴方のプロジェクトのための top-level モジュールをインストールします

   # setup.py
   #!/usr/bin/env python
   
   from setuptools import setup, find_packages
   
   setup(name='src',
         version='0.0.1',
         description='Describe Your Cool Project',
         author='',
         author_email='',
         url='https://github.com/YourSeed',  # REPLACE WITH YOUR OWN GITHUB PROJECT LINK
         install_requires=[
               'pytorch-lightning'
         ],
         packages=find_packages()
         )
   

2. 次のようにプロジェクトをセットアップします :

   /project
       /src
           some_file.py
           /or_a_folder
       setup.py
   

3. root-level パッケージとしてインストールします

   cd /project
   pip install -e .
   

4. それからスクリプトをどこでも呼び出すことができます

   cd /project/src
   python some_file.py --accelerator 'ddp' --gpus 8 
   

 

DP

貴方が Windows を使用している場合、DDP はサポートされません。代わりに DataParallel のための dp を使用できます。DataParallel はマルチプロセス処理の代わりに、マルチスレッドを利用します。それは k GPU に渡りバッチを分割します。つまり、32 のバッチを持ち 2 gpu で DP を使用する場合、各 GPU は 16 サンプルを処理し、その後 root ノードは結果を集計します。

DP の使用は PyTorch と Lightning で奨励されません。DDP を使用してください、これはよりステーブルで少なくとも 3x 高速です。

# dp = DataParallel
trainer = pl.Trainer(gpus=2, accelerator='dp')

trainer.fit(model, train_loader, val_loader)

 

DDP2

ある場合には、サブセットの代わりに、同じマシン上で総てのバッチを使用することは有利です。例えば、自己教師あり学習で、一般的なパフォーマンス・ブーストはネガティブ・サンプルの数を増やすことに由来します。

この場合、DDP2 を利用できます、これは一つのマシンで DP そしてノードに渡り DDP のように動作します。DDP2 は以下を行ないます :

• データのサブセットを各ノードにコピーする。
• 各ノード上のモデルを初期化する。
• DP を使用して forward と backward パスを実行する。
• ノードに渡り勾配を同期する。
• optimizer 更新を適用する。

# ddp2 = DistributedDataParallel + dp
trainer = pl.Trainer(gpus=2, num_nodes=2, accelerator='ddp2')

trainer.fit(model, train_loader, val_loader)

• 2 番目のモードは ddp_spawn です。これは ddp のように動作しますが、貴方のスクリプトを複数回呼び出す代わりに、lightning は GPU 毎にサブプロセスを開始するためにマルチプロセス処理 spawn を使用します。

けれども、このモードを dataloader の num_workers > 0 と混在させることには注意するべきです、何故ならばそれは訓練のボトルネックとなるからです。これは PyTorch の現在の既知の制限です、これは代わりに ddp 実装を何故勧めるかです。

 

mocking ddp

DDP のテストやデバッグは困難であり得ますので、それを容易にするために cpu 上で ddp をシミュレートする分散バックエンドを持ちます。GPU のないマシン上で分散訓練を模倣するために accelerator=”ddp_cpu” を使用するとき num_processes を 1 より大きい数に設定します。デバッグのためにこれが有用である一方で、それはどのようなスピードアップも提供しないことに注意してください、何故ならば単一プロセス Torch はマルチ CPU を既に効率的に利用しているからです。

# Simulate DDP for debugging on your GPU-less laptop
trainer = Trainer(accelerator="ddp_cpu", num_processes=2)

trainer.fit(model, train_loader, val_loader)

 

TPU 上で訓練する

貴方の訓練を加速するためのもう一つのオプションは TPU を使用することです。TPU は深層学習のために特に設計された、Tensor 処理ユニット (Tensor processing unit) です。各 TPU は 8 コアを持ち、そこでは各コアは 128×128 行列乗算のために最適化されています。Google は 8 TPU コアはおよそ 4 V100 GPU のように高速と見積もっています！

TPU pod はその上で多くの TPU をホストしています。現在、TPU pod v2 は 2048 コアを持ちます！Google cloud から完全な pod または (2048 コアのあるサブセットを与える)「スライス」を要求できます。

現時点で、TPU は Google Cloud (GCP), Google Colab と Kaggle 環境で利用可能です。

Lightning は貴方のモデルへのどのようなコード調整もなしで TPU 上の訓練をサポートします。ちょうど GPU を使用するときのように、Lightning は正しいサンプラーを自動的に挿入します – これを貴方自身で行なう必要はありません！

内部的には、lightning は facebook と google XLA チームの連携により開発された XLA フレームワークを使用しています。そして PyTorch の進んだ TPU 採用において彼らの努力を認識することを望みます。

 

tpu_cores

TPU 上で訓練するには、tpu_cores フラグを設定します。

colab or kaggle を使用するとき、許容される値は 1 or 8 コアです。google cloud を使用するとき、8 以上の任意の値が許容されます。

Your effective batch size is the batch size passed into a dataloader times the total number of tpu cores.

# int: train on a single core
trainer = pl.Trainer(tpu_cores=1)

trainer.fit(model, train_loader, val_loader)

# int: train on all cores few cores
trainer = pl.Trainer(tpu_cores=8)

trainer.fit(model, train_loader, val_loader)

リスト [1-8] を渡すことにより、どの TPU コアの上で訓練するかを選択することもできます。これは公式にはサポートされていないユースケースですが、このユーザ体験を改良するために XLA チームと作業しています。

# list: train on a single selected core
trainer = pl.Trainer(tpu_cores=[2])

trainer.fit(model, train_loader, val_loader)

8 コア以上の上で訓練するには、xla_dist スクリプトを使用してこのスクリプトを submit します。

python -m torch_xla.distributed.xla_dist
--tpu=$TPU_POD_NAME
--conda-env=torch-xla-nightly
--env=XLA_USE_BF16=1
-- python your_trainer_file.py

 

進んだ分散訓練

trainer フラグを設定することにより Lightning はマルチ GPU と TPU に渡る分散訓練を out of the box にサポートしますが、必要であればそれはサンプリングが成される方法を制御することを可能にします。

 

replace_sampler_ddp

PyTorch では、マルチノードや GPU 訓練のために torch.nn.DistributedSampler を使用しなければなりません。サンプラーは各 GPU が貴方のデータの適切なパートを見ることを確かなものにします。

# without lightning
def train_dataloader(self):
    dataset = MNIST(...)
    sampler = None
 
    if self.on_tpu:
        sampler = DistributedSampler(dataset)
 
    return DataLoader(dataset, sampler=sampler)

Lightning は必要とされるとき正しいサンプラーを追加しますので、サンプラーを明示的に追加する必要はありません。デフォルトではそれは train サンプラーのために shuffle=True を追加して val/test サンプラーのために shuffle=False を追加します。

この動作をカスタマイズすることを望む場合、replace_sampler_ddp=False を設定して貴方自身の分散サンプラーを追加することができます。

(note: iterable データセットについては、これを自動的には行ないません。)

sampler = torch.utils.data.distributed.DistributedSampler(dataset, shuffle=False)
dataloader = DataLoader(dataset, batch_size=32, sampler=sampler)

trainer = pl.Trainer(gpus=2, num_nodes=2, replace_sampler_ddp=False)

 

prepare_data_per_node

マルチ NODE 訓練を行なうとき、ノードが同じファイルシステムを共有する場合、可能性のある衝突を回避するためにデータを一度以上ダウンロードすることは望みません。

Lightning はマスターノードの root GPU (ie: 単一 GPU だけ) の上で prepare_data フックを自動的に呼び出します。

ノードが同じファイルシステムを共有しないところの幾つかのケースでは、各ノードでデータをダウンロードする必要があります。この場合このフラグを true に設定できてそして lightning は各ノードの root GPU 上でデータをダウンロードします。

このフラグは True にデフォルト設定されています。

trainer = pl.Trainer(gpus=2, num_nodes=2, prepare_data_per_node=False)

trainer.fit(model, train_loader, val_loader)

 

sync_batchnorm

batch norm が GPU/TPU 毎に計算されます。このフラグは総ての GPU に渡る batchnorm 層の間での同期を有効にします。小さいバッチサイズを持つ場合それは推奨されます。

trainer = Trainer(gpus=4, sync_batchnorm=True)

trainer.fit(model, train_loader, val_loader)

以上