Pytorch に基づく分散データの並列トレーニング
はじめに: PyTorch で DistributedDataParallel を使用したマルチ GPU 分散モデル トレーニング
モチベーション
ニューラル ネットワークのトレーニングを高速化する最も簡単な方法は、GPU を使用することです。GPU は、ニューラル ネットワークで一般的なタイプの計算 (行列の乗算と加算) を CPU よりも高速化します。モデルやデータセットが大きくなると、1 つの GPU ではすぐに不足する可能性があります。たとえば、BERT や GPT-2 などの大規模な言語モデルは、数百の GPU でトレーニングされます。マルチ GPU トレーニングを実行するには、モデルとデータを異なる GPU に分割し、トレーニングを調整する方法が必要です。
なぜデータを並列分散するのでしょうか?
多くの人は、ニューラル ネットワーク フレームワークの制御と使いやすさのバランスが最も優れているため、Pytorch で独自の深層学習モデルを実装することを好みます。Pytorch には、モデルとデータを複数の GPU に分割する 2 つの方法があります:nn.DataParallelとnn.DistributedDataParallel。
nn.DataParallel の方が使いやすいです (モデルをラップしてトレーニング スクリプトを実行するだけです)。ただし、1 つのプロセスを使用してモデルの重みを計算し、それを各バッチの各 GPU に分散するため、ネットワークがすぐにボトルネックになり、GPU 使用率が低くなることがよくあります。また、nn.DataParallel ではすべての GPU が同じノード上にある必要があり、混合精度トレーニングに Apex とともに使用することはできません。
したがって、nn.DataParallelとの主な違いはnn.DistributedDataParallel次のように要約できます。
1. DistributedDataParallel はモデルの並列処理をサポートしますが、DataParallel はサポートしません。つまり、モデルが大きすぎて 1 枚のカードのメモリが不十分な場合は、前者のみを使用できます。 ;
2. DataParallel は単一プロセス マルチスレッドであり、単一マシンの状況でのみ使用されますが、DistributedDataParallel はマルチプロセスであり、単一マシンおよび複数マシンの状況に適しており、真に分散トレーニングを実現し
ます。 DistributedDataParallel は、各プロセスが独立した Python インタプリタであるため、GIL 問題を回避し、通信コストが低く、トレーニング速度が速いため、より効率的であり、基本的に DataParallel は放棄されています ; 4. DistributedDataParallel の各プロセスには
、独立したオプティマイザー。独自の更新プロセスを実行しますが、勾配は各プロセスに渡され、内容はすべての実行で同じです。
入手可能な情報が不足している
全体として、Pytorch のドキュメントは完全かつ明確ですが、その使用方法を理解しようとすると、DistributedDataParallelすべての例とチュートリアルが、アクセスできない、不完全な、またはオーバーロードされた無関係な関数の組み合わせであることがわかります。
Pytorch は、AWS を使用した分散トレーニングに関するチュートリアルを提供しており、AWS 側でのセットアップ方法をうまく示しています。ただし、残りの部分は少し混乱しています。何らかの理由で、メトリクスの計算方法を示すのに多くの時間を費やし、その後、モデルをラップしてプロセスを開始する方法の説明に戻るからです。また、それが何を行うのかについても説明されていないためnn.DistributedDataParallel、関連するコード ブロックを理解するのが難しくなります。
Pytorch を使用した分散アプリケーションの作成に関するチュートリアルは、最初のパスに必要な内容よりもはるかに詳細であり、Python マルチプロセッシングのバックグラウンドがない人にはアクセスできません。nn.DistributedDataParallel機能を複製するのに多くの時間を費やします。ただし、その機能の概要や使用方法についての洞察は提供されません ( https://pytorch.org/tutorials/intermediate/ddp_tutorial.html ) 。
分散データ並列処理を開始する方法についてのPytorchチュートリアルもあります。これはセットアップの方法を示していますが、セットアップの目的は説明していません。その後、モデルを GPU 間で分割し、最適化ステップを実行するコードを示しています。残念ながら、書かれたコードは実行できない（関数名が一致しない）ことは確実であり、コードを実行する方法は示されていません。前のチュートリアルと同様に、分散トレーニングがどのように機能するかについての概要は説明しません。
Pytorch が提供する MWE サンプルに最も近いものは、Imagenetトレーニング サンプルです。残念ながら、この例は Pytorch の他のほぼすべての機能も示しているため、分散マルチ GPU トレーニングに何が関連しているかを見つけるのは困難です。
Apex独自のバージョンの Pytorch Imagenet サンプルを提供します。彼らのバージョンの nn.DistributedDataParallel は Pytorch の代替品であり、Pytork の使用方法を学んだ後にのみ役に立ちます。
このチュートリアルでは、内部で何が起こっているのか、そしてそれがどのようにnn.DataParallel異なるのかをうまく説明しています。nn.DataParallelただし、その使用方法に関するコードサンプルはありません。
概要
このチュートリアルは、Pytorch でのニューラル ネットワーク モデルのトレーニングに既に慣れている人を対象としています。全体的なアイデアの概要を説明することから始めます。次に、GPU 上の MNIST を使用したトレーニングの最小限の動作例を示します。この例を変更して、複数の GPU (場合によっては複数のノードにわたって) でトレーニングし、変更を 1 行ずつ説明しました。重要なのは、コードの実行方法についても説明していることです。ボーナスとして、 をApex使用して単純な混合精度分散トレーニングを実行する方法も示します。
全体フレーム図
DistributedDataParallelマルチプロセッシングを使用すると、複数の GPU にわたってモデルが複製され、それぞれが単一のプロセスによって制御されます。(プロセスはコンピューター上で実行される Python のインスタンスです。複数のプロセスを並行して実行することで、複数の CPU コアを持つプロセッサーを活用できます。必要に応じて、各プロセスで複数の GPU を制御することもできますが、これでは明らかに効率が低下します。プロセスごとに 1 つの GPU を使用するよりも遅くなります。複数のワーカー プロセスが各 GPU のデータをフェッチすることも可能ですが、簡略化のために省略します。) GPU はすべて同じノード上に配置することも、複数のノードに分散することもできます。(ノードはすべての CPU と GPU を含む「コンピューター」です。AWS を使用する場合、ノードは EC2 インスタンスです。) 各プロセスは同じタスクを実行し、各プロセスは他のすべてのプロセスと通信します。ネットワーク通信がボトルネックにならないように、プロセス/GPU 間では勾配のみが渡されます。
トレーニング中、各プロセスはディスクから独自のミニバッチをロードし、それらを GPU に渡します。各 GPU には独自のフォワード パスがあり、GPU 全体で勾配が減少します。各層の勾配は前の層に依存しないため、ネットワークのボトルネックをさらに軽減するために、勾配 all-reduce がバックワード パスと同時に計算されます。逆のプロセスの最後には、各ノードに平均勾配が設定され、モデルの重みが同期した状態に保たれます。
これらはすべて、同期および通信するために複数のプロセス (場合によっては複数のノード上) を必要とします。Pytorch はdistributed.init_process_groupその機能を通じてこれを実現します。この関数は、すべてのプロセスが同期できるようにプロセス 0 を見つける場所と、予想されるプロセスの合計数を知る必要があります。個々のプロセスは、プロセスの総数とプロセス内でのランク、およびどの GPU が使用されているかを知る必要もあります。プロセスの総数を指すのがworld size一般的です。最後に、バッチが重複しないように、各プロセスはデータのどの部分を処理するかを認識する必要があります。Pytorch はnn.utils.data.DistributedSamplerこれを実現するために、トレーニング データが重複しないようにプロセスごとにデータを分割します。
DDP の内部メカニズムの詳細については、公式ドキュメントを参照してください: DISTRIBUTED DATA PARALLEL
説明付きの最小限のデモ例
これを行う方法を示すために、MNIST でトレーニングするサンプルを作成し、それを複数のノードにわたる複数の GPU で実行するように変更し、最後に混合精度トレーニングも可能にします。
マルチプロセッシングなし
まず、必要な依存関係をインポートします。
import os
import argparse
import torch
import torch.nn as nn
import torch.distributed as dist
import torch.multiprocessing as mp
import torchvision
import torchvision.transforms as transforms
from datetime import datetime
from apex.parallel import DistributedDataParallel as DDP
from apex import amp
MNIST を予測するための非常に単純な畳み込みモデルを定義します。
class ConvNet(nn.Module):
def __init__(self, num_classes=10):
super(ConvNet, self).__init__()
self.layer1 = nn.Sequential(
nn.Conv2d(1, 16, kernel_size=5, stride=1, padding=2),
nn.BatchNorm2d(16),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2)
)
self.layer2 = nn.Sequential(
nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),
nn.BatchNorm2d(32),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2)
)
self.fc = nn.Linear(7 * 7 * 32, num_classes)
def forward(self, x):
out = self.layer1(x)
out = self.layer2(out)
out = out.reshape(out.size(0), -1)
out = self.fc(out)
return out
トレーニングのプロセスは次のとおりです。
def train(gpu, args):
torch.manual_seed(0)
model = ConvNet()
torch.cuda.set_device(gpu)
model.cuda(gpu)
# model = nn.DataParallel(model, device_ids=device_ids)
# model = model.cuda(device=gpu)
batch_size = 100
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda(gpu)
optimizer = torch.optim.SGD(model.parameters(), 1e-4)
# Data loading code
train_dataset = torchvision.datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor(),
download=True)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True,
num_workers=0, pin_memory=True)
start = datetime.now()
total_step = len(train_loader)
for epoch in range(args.epochs):
for i, (images, labels) in enumerate(train_loader):
images = images.cuda(non_blocking=True)
labels = labels.cuda(non_blocking=True)
# Forward pass
outputs = model(images)
loss = criterion(outputs, labels)
# Backward and optimize
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (i + 1) % 100 == 0 and gpu == 0:
print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(epoch + 1, args.epochs, i + 1, total_step,
loss.item()))
if gpu == 0:
print("Training complete in: " + str(datetime.now() - start))
main() 関数はいくつかのパラメータを受け取り、トレーニング関数を実行します。
def main():
parser = argparse.ArgumentParser()
parser.add_argument("-n", "--nodes", default=1, type=int, metavar='N')
parser.add_argument('-g', '--gpus', default=1, type=int, help='number of gpus per node')
parser.add_argument('-nr', '--nr', default=0, type=int, help='ranking within the nodes')
parser.add_argument('--epochs', default=2, type=int, metavar='N', help='number of total epochs to run')
args = parser.parse_args()
train(0, args)
最後に、main() 関数が呼び出されることを確認します。
if __name__ == '__main__':
main()
ターミナルを開いて「 」と入力することでこのコードを実行できpython src/mnist.py-n 1-g 1-nr 0、単一ノードの単一 GPU でトレーニングされます。
マルチプロセッシングを有効にする
これをマルチプロセッシングで行うには、各 GPU のプロセスを開始するスクリプトが必要です。各プロセスは、どの GPU を使用するか、および実行中のすべてのプロセスの中でのそのランクを知る必要があります。スクリプトは各ノードで実行する必要があります。
各機能の変更点を見てみましょう。新しいコードは見つけやすいように分離されています。
def main():
parser = argparse.ArgumentParser()
parser.add_argument("-n", "--nodes", default=1, type=int, metavar='N')
parser.add_argument('-g', '--gpus', default=1, type=int, help='number of gpus per node')
parser.add_argument('-nr', '--nr', default=0, type=int, help='ranking within the nodes')
parser.add_argument('--epochs', default=2, type=int, metavar='N', help='number of total epochs to run')
args = parser.parse_args()
#########################################################
args.world_size = args.gpus * args.nodes
os.environ['MASTER_ADDR'] = '172.20.109.105'
os.environ['MASTER_PORT'] = '8888'
mp.spawn(train, nprocs=args.gpus, args=(args,))
#########################################################
# train(0, args)
の：
args.nodesノードの総数を表します。args.gpusノードあたりの GPU の合計数を示します (ノードあたりの GPU の数は同じです)。args.nr全ノードのうち、現在のノードのシリアル番号を示します。
world_sizeノードの総数とノードごとの GPU の数に従って、実行されるプロセスの総数を計算できます。すべてのプロセスはプロセス 0 の IP アドレスとポートを知っている必要があるため、すべてのプロセスは通常、プロセス 0 はマスター プロセスと呼ばれ、たとえばプロセス 0 で情報を出力したり、モデルを保存したりします。
PyTorch は、ノード上でノードのすべてのプロセスを開始しmp.spawn、各プロセスを実行します。ここで、i は 0 から-1 です。合計のプロセスが存在するように、各ノードで main() 関数を実行することを忘れないでください。train(i, args)args.gpusargs.nodes*args.gpus=args.world_size
もう一度、トレーニング関数を変更します。
def train(gpu, args):
############################################################
rank = args.nr * args.gpus + gpu
dist.init_process_group(backend='nccl', init_method='env://', world_size=args.world_size, rank=rank)
############################################################
torch.manual_seed(0)
model = ConvNet()
torch.cuda.set_device(gpu)
model.cuda(gpu)
batch_size = 100
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda(gpu)
optimizer = torch.optim.SGD(model.parameters(), 1e-4)
############################################################
# Wrap the model
model = nn.parallel.DistributedDataParallel(model, device_ids=[gpu])
############################################################
# Data loading code
train_dataset = torchvision.datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor(),
download=True)
############################################################
train_sampler = torch.utils.data.distributed.DistributedSampler(dataset=train_dataset, num_replicas=args.world_size,
rank=rank)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=False,
num_workers=0, pin_memory=True, sampler=train_sampler)
############################################################
start = datetime.now()
total_step = len(train_loader)
for epoch in range(args.epochs):
for i, (images, labels) in enumerate(train_loader):
images = images.cuda(non_blocking=True)
labels = labels.cuda(non_blocking=True)
# Forward pass
outputs = model(images)
loss = criterion(outputs, labels)
# Backward and optimize
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (i + 1) % 100 == 0 and gpu == 0:
print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(epoch + 1, args.epochs, i + 1, total_step,
loss.item()))
if gpu == 0:
print("Training complete in: " + str(datetime.now() - start))
ここでは、まず現在のプログラム番号を計算してrank = args.nr * args.gpus + gpuから、dist.init_process_group分散環境を初期化します。
backendパラメータは、、 、を含む通信バックエンドを指定します。ここで選択されます。Nvidia が提供する公式のマルチカード通信フレームワークであり、比較的効率的です。これはハイ パフォーマンス コンピューティングの一般的な通信プロトコルでもありますが、MPI 実装フレームワークを自分でインストールする必要があります。通信バックエンドが組み込まれていますが、効率が十分ではありません。mpiglooncclncclmpiOpenMPIglooinit_method最初にプロセスの同期を完了するための初期化方法を指します。ここで設定するのは、env://環境変数の初期化メソッドを指します。環境変数には 4 つのパラメーターを設定する必要があります:MASTER_PORT、MASTER_ADDR、WORLD_SIZE、RANK、最初の 2 つのパラメーターはすでに設定されています。後の 2 つのパラメーターは、dist.init_process_group関数world_size中和rankパラメーターを通じて構成することもできます。
他の初期化方法には、共有ファイル システムやTCP などがあります。たとえば、TCP が初期化方法として使用される場合はinit_method='tcp://10.1.1.20:23456'、マスターの IP アドレスとポートを指定する必要があります。この呼び出しはブロックされているため、すべてのプロセスが同期するまで待機する必要があり、いずれかのプロセスが失敗すると失敗することに注意してください。DistributedDataParallelモデル側では、元のモデルをラップし、処理のためにモデルを GPU にコピーするだけで済み、バックグラウンドで勾配 All-Reduce 操作がサポートされます。- データ側では、プロセスごとにデータを分割するためにこれを使用します。これを使用する
nn.utils.data.DistributedSampler必要があるのは、 でのみです。トレーニング サイクル プロセスの各エポックの開始時に呼び出す必要があることに注意してください(主に、プロセスの分割を確実にするため)。各エポックは異なります）他のトレーニング コードは変更されません。dataloadersamplertrain_sampler.set_epoch(epoch)
最後に、コードを実行できます。たとえば、4 つのノードがあり、各ノードに 8 枚のグラフィック カードがある場合、4 つのノードの端末上でコードを実行する必要があります。
python src/mnist-distributed.py -n 4 -g 8 -nr i
たとえば、ノード 0 で実行します。
python src/mnist-distributed.py -n 4 -g 8 -nr 0
つまり、各ノードでこのスクリプトを実行すると、トレーニングが開始される前に相互に同期しているargs.gpusプロセスを開始するように指示されます。
現時点で効果的なのは、 BN を使用したモデルのbatch_size場合batch_size_per_gpu * world_size、同期 BN を使用してより良い結果を得ることができることです。
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
上記では、分散トレーニング プロセスについて説明しましたが、これは評価またはテスト プロセスにも適用できます。たとえば、データを予測のために複数の異なるプロセスに分割することで、予測プロセスを高速化できます。実装コードは上記の処理と全く同じですが、各プロセスはデータ内容の一部しか計算していないため、ある指標を計算したい場合は各プロセスの統計結果からAll-Reduceする必要があります。例えば、分類精度を計算したい場合、データ総数と各プロセスの分類正解数をカウントし、それらを集計することができます。
ここで言及すべきことの 1 つは、分散環境を初期化するときに、デフォルトの分散プロセス グループdist.init_process_group(分散プロセス グループ)が実際に確立され、このグループは Pytorchパッケージも初期化するということです。このようにして、直接使用できる API で分散基本操作を実行できます。具体的な実装は次のとおりです。torch.distributedtorch.distributed
# define tensor on GPU, count and total is the result at each GPU
t = torch.tensor([count, total], dtype=torch.float64, device='cuda')
dist.barrier() # synchronizes all processes
dist.all_reduce(t, op=torch.distributed.ReduceOp.SUM,) # Reduces the tensor data across all machines in such a way that all get the final result.
t = t.tolist()
all_count = int(t[0])
all_total = int(t[1])
acc = all_count / all_total
分散トレーニングの開始方法
上記のプロセスでは、PyTorchtorch.multiprocessingパッケージ ( Multiprocessing パッケージ - torch.multiprocessing ) を使用して分散トレーニングを開始しています。現在、公式のImageNet トレーニング サンプルがこの方法で使用されており、detectron2 ライブラリも次の方法で開始されています: https:/ /github.com/facebookresearch/detectron2/blob/main/detectron2/engine/launch.py。
スタートアップを使用する場合は、入力トレーニングがこの形式である必要があるtorch.multiprocessing.spawnことに注意する必要があります。最初のパラメータ i は現在のノードのプロセス番号を指し、このパラメータは実際には現在のノードのいわゆるトレーニング プロセスとして機能します。番号, 上記のランクは実際にはグローバル プログラム番号です. このパラメータに従って各プロセスで使用されるデバイス デバイスを設定する必要があるため、このパラメータは非常に重要です. 一般に、これは使用される GPU 番号と直接見なされ、設定は次のとおりです。functionfn(i,*args)local_ranklocal_ranklocal_rank
torch.cuda.set_device(args.local_rank) # before your code runs
# set DDP
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank)
# 或者
with torch.cuda.device(args.local_rank):
# your code to run
の採用に加えて、プログラム (分散通信パッケージ - torch.distributedmp.spawn ) を起動するために使用することもできます。これは、より一般的な起動方法です。たとえば、単一マシンのマルチカード トレーニングの場合、起動方法は次のとおりです。torch.distributed.launch
python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE
YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 and all other arguments of your training script)
その中にはNUM_GPUS_YOU_HAVE、GPU の総量YOUR_TRAINING_SCRIPT.pyと、基本的には上記と同じですが、一部の環境変数 () が起動時に自動的に設定さtorch.distributed.launchれる点がhttps://github.com/pytorch/pytorch/blob/master/torch/distributed/run.py#L211異なります。環境変数から直接:RANKWORLD_SIZE
rank = int(os.environ["RANK"])
world_size = int(os.environ['WORLD_SIZE'])
local_rank取得するには 2 つの方法があります: 1
) 1 つはコマンド ライン パラメータをトレーニング スクリプトに追加する方法です。これはプログラムの開始時に自動的に割り当てられます。
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int)
args = parser.parse_args()
local_rank = args.local_rank
2) もう 1 つの方法は、torch.distributed.launchplus を起動することです--use_env=True。この場合、LOCAL_RANKこの環境変数が設定され、環境変数から取得できますlocal_rank。
"""
python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE --use_env=True
YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 and all other arguments of your training script)
"""
import os
local_rank = int(os.environ["LOCAL_RANK"])
2 つのノードなどのマルチマシン マルチカード トレーニングの場合、起動コマンドは次のとおりです。
# Node 1: (IP: 192.168.1.1, and has a free port: 1234)
python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE --nnodes=2 --node_rank=0 --master_addr="192.168.1.1"
--master_port=1234 YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 and all other arguments of your training script)
# Node 2
python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE --nnodes=2 --node_rank=1 --master_addr="192.168.1.1"
--master_port=1234 YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 and all other arguments of your training script)
ここ
--nnodes受信ノードの数を示します--node_rank着信ノードの番号を示しますworld_size=nnodes*nproc_per_node。
ただし、PyTorch の最新バージョンでは、代わりにtorchruntorch.distributed.launchが起動されました。torchrunと の使い方はtorch.distributed.launch基本的に同じですが、コマンドは放棄され--use_env、local_rank環境変数に直接設定されます。最新バージョンでは起動方法をtorchvision使用しています。詳細については、main · pytorch/vision のvision/references/classificationを参照してください。torchrun
混合精度トレーニング (apex を使用)
Apexをインストールします。
git clone https://github.com/NVIDIA/apex
cd apex
# if pip >= 23.1 (ref: https://pip.pypa.io/en/stable/news/#v23-1) which supports multiple `--config-settings` with the same key...
pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" ./
# otherwise
pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --global-option="--cpp_ext" --global-option="--cuda_ext" ./
Apex 公式ドキュメント: Apex (PyTorch 拡張機能)
混合精度トレーニング (浮動小数点 (FP32) トレーニングと半精度 (FP16) トレーニングの組み合わせ) では、より大きなバッチ サイズを使用し、NVIDIA Tensor コアを活用して計算を高速化できます。AWS p3 インスタンスは、Tensor コアを備えた NVIDIA Tesla V100 GPU を使用します。NVIDIA の apex を混合精度トレーニングに使用するのは非常に簡単で、コードの一部を変更するだけです。
def train(gpu, args):
############################################################
rank = args.nr * args.gpus + gpu
dist.init_process_group(backend='nccl', init_method='env://', world_size=args.world_size, rank=rank)
############################################################
torch.manual_seed(0)
model = ConvNet()
torch.cuda.set_device(gpu)
model.cuda(gpu)
batch_size = 100
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda(gpu)
optimizer = torch.optim.SGD(model.parameters(), 1e-4)
############################################################
# Wrap the model
model, optimizer = amp.initialize(model, optimizer, opt_level='O2')
model = DDP(model)
############################################################
# Data loading code
train_dataset = torchvision.datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor(),
download=True)
############################################################
# train_sampler = torch.utils.data.distributed.DistributedSampler(dataset=train_dataset, num_replicas=args.world_size,
# rank=rank)
# train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=False,
# num_workers=0, pin_memory=True, sampler=train_sampler)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True,
num_workers=0, pin_memory=True)
############################################################
start = datetime.now()
total_step = len(train_loader)
for epoch in range(args.epochs):
for i, (images, labels) in enumerate(train_loader):
images = images.cuda(non_blocking=True)
labels = labels.cuda(non_blocking=True)
# Forward pass
outputs = model(images)
loss = criterion(outputs, labels)
# Backward and optimize
optimizer.zero_grad()
############################################################
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
############################################################
optimizer.step()
if (i + 1) % 100 == 0 and gpu == 0:
print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(epoch + 1, args.epochs, i + 1, total_step,
loss.item()))
if gpu == 0:
print("Training complete in: " + str(datetime.now() - start))
実際には次の 2 つの変更があります。
- 1 つ目は、混合精度トレーニングを
amp.initializeパッケージ化してサポートするためにmodel使用します。これは、最適化レベルを指します。O0 (すべての浮動小数点を使用) または O3 (全体で半精度を使用) の場合、実際の混合精度ではありませんが、可能です。モデル効果と速度のベースラインを決定するために使用されます。O1 と O2 は混合精度の 2 つの設定です。混合精度トレーニング用に 1 つを選択できます。詳細は、 Apex ドキュメントを参照してください。optimizeropt_level - もう 1 つのポイントは、勾配に従ってパラメーターを更新する前に、勾配のアンダーフローを防ぐために amp.scale_loss によって勾配をスケーリングする必要があることです。また、nn.DistributedDataParallel を apex.Parallel.DistributedDataParallel に置き換えることもできます。
はい、これらのコードの最初の文字は大文字の「O」で、2 番目の文字は数字です。はい、ゼロを代入すると、不可解なエラー メッセージが表示されます。
apex.parallel.distributedDataParallelはい、nn.distributedDataParallear交換品です。Apex ではプロセスごとに 1 つの GPU しか許可されないため、GPU を指定する必要がなくなりました。また、モデルを GPU に移動する前にスクリプトが呼び出されることも前提としていますtorch.cuda.set_device(local_rank)。
混合精度のトレーニングでは、勾配のアンダーフローを防ぐために損失をスケーリングする必要があります。Apex はこれを自動的に実行します。
このスクリプトは、分散トレーニング スクリプトと同じように機能します。
python without_multiprocessing.py -n 1 -g 4 -nr 0
さらに、新しいバージョンの PyTorch には混合精度トレーニングが組み込まれています。詳細については、「AUTOMATIC MIXED PRECISION PACKAGE - TORCH.AMPリンクの説明を追加」を参照してください。さらに、PyTorch の公式分散実装は現在比較的完成しており、そのパフォーマンスと効果は良好です。代替ソリューションとしては、horovodPyTorch だけでなく、TensorFlow および MXNet フレームワークもサポートすることです。実装は比較的簡単で、速度も良好です。