本文已参与「新人创作礼」活动，一起开启掘金创作之路。

相关教程

MistGPU地址 PyCharm连接MistGPU教程

MistGPU使用注意事项

ifconfig命令查看ip地址报错

报错：zsh: command not found: ifconfig 原因：在服务器上第一次使用该命令需要先安装net-tools 解决办法：sudo apt install net-tools

sudo apt install net-tools安装网络包报错

报错：unable to locate package net-tools 解决办法：sudo apt-get update

在MistGPU上训练

ifconfig查看服务器ip地址 并将参数配置中的服务器ip地址改为172.17.0.2

# mnist-cpu-distributed.py
import os
from datetime import datetime
import argparse
import torch.multiprocessing as mp
import torchvision
import torchvision.transforms as transforms
import torch
import torch.nn as nn
import torch.distributed as dist

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 定义一个简单的CNN模型处理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):
    rank = args.nr * args.gpus + gpu # 首先计算出当前进程序号：rank = args.nr * args.gpus + gpu

    dist.init_process_group( # 通过dist.init_process_group初始化分布式环境
        backend='nccl', # backend参数指定通信后端，包括mpi, gloo, nccl，这里选择nccl，这是Nvidia提供的官方多卡通信框架，相对比较高效。
        init_method='env://',
        # init_method指的是如何初始化，以完成刚开始的进程同步；这里我们设置的是env://，
        # 指的是环境变量初始化方式，需要在环境变量中配置4个参数：MASTER_PORT，MASTER_ADDR，WORLD_SIZE，RANK，
        # 前面两个参数我们已经配置，后面两个参数也可以通过dist.init_process_group函数中world_size和rank参数配置。
        # 其它的初始化方式还包括共享文件系统以及TCP，比如init_method='tcp://10.1.1.20:23456'，其实也是要提供master的IP地址和端口。
        # 注意这个调用是阻塞的，必须等待所有进程来同步，如果任何一个进程出错，就会失败。
        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
    # 对于模型侧，我们只需要用DistributedDataParallel包装一下原来的model即可，在背后它会支持梯度的All-Reduce操作。
    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
    )

    # 对于数据侧，我们nn.utils.data.DistributedSampler来给各个进程切分数据，只需要在dataloader中使用这个sampler就好，
    # 值得注意的一点是你要训练循环过程的每个epoch开始时调用train_sampler.set_epoch(epoch)，（
    # 主要是为了保证每个epoch的划分是不同的）其它的训练代码都保持不变。
    train_sampler = torch.utils.data.distributed.DistributedSampler(
        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.to(device)
            labels = labels.to(device)
            # 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()
    args.world_size = args.gpus * args.nodes
    os.environ['MASTER_ADDR'] = '172.17.0.2' # 服务器ip地址
    os.environ['MASTER_PORT'] = '8888'
    mp.spawn(train, nprocs=args.gpus, args=(args,))

# 通过启动主函数来开始训练
if __name__ == '__main__':
    main()

命令行执行代码

python src/mnist-distributed.py -n 4 -g 8 -nr i