"""
从torch.distributed.launch的源码,可以看出launch实际上主要完成的工作:
1.参数定义与传递。解析环境变量,并将变量传递到子进程中。
2.起多进程。调用subprocess.Popen启动多进程。
用launch方式需要注意的位置:
需要添加一个解析local_rank的参数
parser.add_argument("--local_rank", type=int)
dist初始化的方式int_method取env
dist.init_process_group("gloo", init_method='env://')
DDP的设备都需要指定
local_rank net=torch.nn.parallel.DistributedDataParallel(net, device_ids=[args.local_rank], output_device=args.local_rank)
"""
# 下面为实际例子
import torch
import torchvision
# 改动1
import torch.utils.data.distributed
import argparse
# 改动2
import torch.distributed as dist
from torchvision import transforms
parser = argparse.ArgumentParser()
# 改动3
parser.add_argument("--local_rank", type=int)
args = parser.parse_args()
# 改动4
torch.cuda.set_device(args.local_rank)
def main():
# 改动5
dist.init_process_group("nccl", init_method='env://')
trans = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
data_set = torchvision.datasets.MNIST('~/DATA/',
train=True,
transform=trans,
target_transform=None,
download=True)
train_sampler = torch.utils.data.distributed.DistributedSampler(data_set)
data_loader_train = torch.utils.data.DataLoader(dataset=data_set,
batch_size=256,
sampler=train_sampler,
num_workers=16,
pin_memory=True)
net = torchvision.models.resnet101(num_classes=10)
net.conv1 = torch.nn.Conv2d(1, 64, (7, 7), (2, 2), (3, 3), bias=False)
net = net.cuda()
# DDP输出方式修改
# 改动6
net = torch.nn.parallel.DistributedDataParallel(net, device_ids=[args.local_rank],
output_device=args.local_rank)
criterion = torch.nn.CrossEntropyLoss()
opt = torch.optim.Adam(net.parameters(), lr=0.001)
for epoch in range(1):
for i, data in enumerate(data_loader_train):
images, labels = data
# 改动7
images = images.to(args.local_rank, non_blocking=True)
labels = labels.to(args.local_rank, non_blocking=True)
opt.zero_grad()
outputs = net(images)
loss = criterion(outputs, labels)
loss.backward()
opt.step()
if i % 10 == 0:
print("loss: {}".format(loss.item()))
if __name__ == '__main__':
main()
"""
假设一共有两台机器(节点1和节点2),每个节点上有8张卡, 节点1的IP地址为192.168.1.1 占用的端口12355(端口可以更换),启动的方式如下:
>>> #节点1
>>>python -m torch.distributed.launch --nproc_per_node=8 --nnodes=2 --node_rank=0 --master_addr="192.168.1.1" --master_port=12355 MNIST.py
>>> #节点2
>>>python -m torch.distributed.launch --nproc_per_node=8 --nnodes=2 --node_rank=1 --master_addr="192.168.1.1" --master_port=12355 MNIST.py
如果只要启动一台机器,只需要将nnodes数量改为1,就是单机多卡的另一种方式。设置如下: python -m torch.distributed.launch --nproc_per_node=8 --nnodes=1 --node_rank=0 --master_addr=localhost --master_port=12355 MNIST.py
"""