为什么需要分布式

众所周知，深度神经网络发展到现阶段，离不开GPU和数据。经过这么多年的积累，GPU的计算能力越来越强，数据也积累的越来越多，大家会发现在现有的单机单卡或者单机多卡上很难高效地复现模型，甚至对于有些新的数据集来讲，单机训练简直就是噩梦。

DatasetImages MS COCO115,000 Open Image dataset v4 1,740,000

为什么单机8卡也会是噩梦呢？我们拿COCO和Google最近Release出来的Open Image dataset v4来做比较，训练一个resnet152的检测模型，在COCO上大概需要40个小时，而在OIDV4上大概需要40天，这还是在各种超参数正确的情况下，如果加上调试的时间，可能一个模型调完就该过年了吧。

所以这个时候我们需要分布式。

Pytorch 分布式简介

PyTorch 1.0稳定版终于正式发布了！新版本增加了JIT编译器、全新的分布式包、C++ 前端，以及Torch Hub等新功能，支持AWS、谷歌云、微软Azure等云平台。

torch.distributed软件包和torch.nn.parallel.DistributedDataParallel模块由全新的、重新设计的分布式库提供支持。

新的库的主要亮点有：

新的 torch.distributed 是性能驱动的，并且对所有后端 (Gloo，NCCL 和 MPI) 完全异步操作
显着的分布式数据并行性能改进，尤其适用于网络较慢的主机，如基于以太网的主机
为torch.distributedpackage中的所有分布式集合操作添加异步支持
在Gloo后端添加以下CPU操作：send，recv，reduce，all_gather，gather，scatter
在NCCL后端添加barrier操作
为NCCL后端添加new_group支持

1.0的多机多卡的计算模型并没有采用主流的Parameter Server结构，而是直接用了Uber Horovod的形式，也是百度开源的RingAllReduce算法。

采用PS计算模型的分布式，通常会遇到网络的问题，随着worker数量的增加，其加速比会迅速的恶化，例如resnet50这样的模型，目前的TF在10几台机器的时候，加速比已经开始恶化的不可接受了。因此，经常要上RDMA、InfiniBand等技术，并且还带来了一波网卡的升级，有些大厂直接上了100GBs的网卡，有钱任性。而Uber的Horovod，采用的RingAllReduce的计算方案，其特点是网络通信量不随着worker（GPU）的增加而增加，是一个恒定值。简单看下图理解下，GPU 集群被组织成一个逻辑环，每个GPU有一个左邻居、一个右邻居，每个GPU只从左邻居接受数据、并发送数据给右邻居。即每次梯度每个gpu只获得部分梯度更新，等一个完整的Ring完成，每个GPU都获得了完整的参数。

这里引入了一个新的函数model = torch.nn.parallel.DistributedDataParallel(model)为的就是支持分布式模式

不同于原来在multiprocessing中的model = torch.nn.DataParallel(model,device_ids=[0,1,2,3]).cuda()函数，这个函数只是实现了在单机上的多GPU训练，根据官方文档的说法，甚至在单机多卡的模式下，新函数表现也会优于这个旧函数。

这里要提到两个问题：

每个进程都有自己的Optimizer同时每个迭代中都进行完整的优化步骤，虽然这可能看起来是多余的，但由于梯度已经聚集在一起并跨进程平均，因此对于每个进程都是相同的，这意味着不需要参数广播步骤，从而减少了在节点之间传输张量tensor所花费的时间。
另外一个问题是Python解释器的，每个进程都包含一个独立的Python解释器，消除了来自单个Python进程中的多个执行线程，模型副本或GPU的额外解释器开销和“GIL-thrashing”。这对于大量使用Python运行时的模型尤其重要。

初始化

torch.distributed.init_process_group(backend, init_method='env://', **kwargs)

参数说明：

backend(str): 后端选择，包括上面那几种 gloo,nccl,mpi
init_method(str，optional): 用来初始化包的URL, 我理解是一个用来做并发控制的共享方式
world_size(int, optional): 参与这个工作的进程数
rank(int,optional): 当前进程的rank
group_name(str,optional): 用来标记这组进程名的

Backends

Backends that come with PyTorch

PyTorch distributed currently only supports Linux. By default, the Gloo and NCCL backends are built and included in PyTorch distributed (NCCL only when building with CUDA). MPI is an optional backend that can only be included if you build PyTorch from source. (e.g. building PyTorch on a host that has MPI installed.)

Which backend to use?

In the past, we were often asked: “which backend should I use?”.

Rule of thumb
- Use the NCCL backend for distributed GPU training
- Use the Gloo backend for distributed CPU training.

init_method分析

‘init_method’支持三种方式：

1. TCP initialization

tcp:// IP组播（要求所有进程都在同一个网络中）比较好理解, 以TCP协议的方式进行不同分布式进程之间的数据交流，需要设置一个端口，不同进程之间公用这一个端口，并且设置host的级别和host的数量。设计两个参数rank和world_size。其中rank为host的编号，默认0为主机，端口应该位于该主机上。world_size为分布式主机的个数。

用该方式，运行上面的代码可以使用如下指令：

在主机01上：

python mnsit.py --init-method tcp://10.172.1.2:22225 --rank 0 --world-size 2

在主机02上：

python mnsit.py --init-method tcp://10.172.1.2:22225 --rank 1 --world-size 2

这里没有设置backend参数，所以默认是gloo。22225是端口号，用一个没有没占用的就行。这两句指令的先后顺序没有要求，只有两条指令都输入，程序才会运行起来。

2. Shared file-system initialization

file:// 共享文件系统（要求所有进程可以访问单个文件系统）有共享文件系统可以选择

提供的第二种方式是文件共享，机器有共享的文件系统，故可以采用这种方式，也避免了基于TCP的网络传输。这里使用方式是使用绝对路径在指定一个共享文件系统下不存在的文件。

在主机01上：

python mnsit.py --init-method file://PathToShareFile/MultiNode --rank 0 --world-size 2

在主机02上：

python mnsit.py --init-method file://PathToShareFile/MultiNode --rank 1 --world-size 2

这里相比于TCP的方式麻烦一点的是运行完一次必须更换共享的文件名，或者删除之前的共享文件，不然第二次运行会报错。

3. Environment variable initialization

env:// 环境变量（需要您手动分配等级并知道所有进程可访问节点的地址）默认是这个

MASTER_PORT - required; has to be a free port on machine with rank 0
MASTER_ADDR - required (except for rank 0); address of rank 0 node
WORLD_SIZE - required; can be set either here, or in a call to init function
RANK - required; can be set either here, or in a call to init function

但是前两个并没有在init_process_group的参数里。这里有一个官方文档中的用例：

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)

需要注意的点

一定要加rank, world_size参数
train_dataset最好不要用自己写的sampler，否则还需要再实现一遍分布式的数据划分方式

Dataloader中的参数

如果你的选项刚好是最坏情况，优化这个有可能达到2倍左右的性能提升（经验值哈）

解释一下DataLoader中其中两个参数：

num_worker：
数据集加载的时候，控制用于同时加载数据的线程数（默认为0，即在主线程读取）存在最优值，你会看到运行的时候pytorch会新建恰等于这个值的数据读取线程，我猜，线程多于必要的时候，数据读取线程返回到主线程反而会因为线程间通信减慢数据。因此大了不好小了也不好。建议把模型，loss，优化器全注释了只跑一下数据流速度，确定最优值
pin_memory：
是否提前申请CUDA内存（默认为False，但有说法除非数据集很小，否则在N卡上推荐总是打开）在MNIST这样的小数据集上好像是关闭比较好，到底多小算小说不清楚，建议自己试一下。

pin_memory就是锁页内存，创建DataLoader时，设置pin_memory=True，则意味着生成的Tensor数据最开始是属于内存中的锁页内存，这样将内存的Tensor转义到GPU的显存就会更快一些。
主机中的内存，有两种存在方式，一是锁页，二是不锁页，锁页内存存放的内容在任何情况下都不会与主机的虚拟内存进行交换（注：虚拟内存就是硬盘），而不锁页内存在主机内存不足时，数据会存放在虚拟内存中。显卡中的显存全部是锁页内存,当计算机的内存充足的时候，可以设置pin_memory=True。当系统卡住，或者交换内存使用过多的时候，设置pin_memory=False。因为pin_memory与电脑硬件性能有关，pytorch开发者不能确保每一个炼丹玩家都有高端设备，因此pin_memory默认为False。

如果机子的内存比较大，建议开启pin_memory=Ture，如果开启后发现有卡顿现象或者内存占用过高，此时建议关闭。

总之官方的默认值很有可能不是最好的，建议自己多试试。在MNIST这样的小数据集上，pin_memory关闭比较好。而且，num_workers需要调节，除了默认情况外，最快和最慢是有一定差距的，建议在自己的代码上只跑数据读取这一块，确定这两个参数的最优值。

分布式 Hello World

启动辅助工具 Launch utility

torch.distributed.launch 例子

torch.distributed包提供了一个启动实用程序torch.distributed.launch，此帮助程序可用于为每个节点启动多个进程以进行分布式训练，它在每个训练节点上产生多个分布式训练进程。

这个工具可以用作CPU或者GPU，如果被用于GPU，每个GPU产生一个进程Process

该工具既可以用来做单节点多GPU训练，也可用于多节点多GPU训练。如果是单节点多GPU，将会在单个GPU上运行一个分布式进程，据称可以非常好地改进单节点训练性能。如果用于多节点分布式训练，则通过在每个节点上产生多个进程来获得更好的多节点分布式训练性能。如果有Infiniband接口则加速比会更高。

在单节点分布式训练或多节点分布式训练的两种情况下，该工具将为每个节点启动给定数量的进程（--nproc_per_node）。如果用于GPU培训，则此数字需要小于或等于当前系统上的GPU数量（nproc_per_node），并且每个进程将在从GPU 0到GPU（nproc_per_node - 1）的单个GPU上运行。

Single-Node multi-process distributed training

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)

Multi-Node multi-process distributed training: (e.g. two nodes)

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)

需要注意的地方：

后端最好用“NCCL”，才能获取最好的分布式性能
训练代码必须从命令行解析--local_rank=LOCAL_PROCESS_RANK

import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int)
args = parser.parse_args()

torch.cuda.set_device(arg.local_rank)

torch.distributed初始化方式

torch.distributed.init_process_group(backend='nccl',
                                     init_method='env://')

model

model = torch.nn.parallel.DistributedDataParallel(model,
                                                  device_ids=[arg.local_rank],
                                                  output_device=arg.local_rank)

其他地方一般就不用修改了，真的顺滑

我们的训练代码中这样写(省略多余代码，只保留核心代码)：

import torch.distributed as dist
# 这个参数是torch.distributed.launch传递过来的，我们设置位置参数来接受，local_rank代表当前程序进程使用的GPU标号
parser.add_argument("--local_rank", type=int, default=0) 

def synchronize():
    """
    Helper function to synchronize (barrier) among all processes when
    using distributed training
    """
    if not dist.is_available():
        return
    if not dist.is_initialized():
        return
    world_size = dist.get_world_size()
    if world_size == 1:
        return
    dist.barrier()


## WORLD_SIZE 由torch.distributed.launch.py产生 具体数值为 nproc_per_node*node(主机数，这里为1)
num_gpus = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1

is_distributed = num_gpus > 1

if is_distributed:
    torch.cuda.set_device(args.local_rank)  # 这里设定每一个进程使用的GPU是一定的
    torch.distributed.init_process_group(
        backend="nccl", init_method="env://"
    )
    synchronize()

# 将模型移至到DistributedDataParallel中，此时就可以进行训练了
if is_distributed:
    model = torch.nn.parallel.DistributedDataParallel(
        model, device_ids=[args.local_rank], output_device=args.local_rank,
        # this should be removed if we update BatchNorm stats
        broadcast_buffers=False,
    )

# 注意，在测试的时候需要执行 model = model.module

WRITING DISTRIBUTED APPLICATIONS WITH PYTORCH

https://github.com/pytorch/examples/tree/master/imagenet
这里，常规的操作就不多叙述了，主要讲一下和分布式相关的代码部分。

parser.add_argument('--world-size', default=2, type=int, help='number of distributed processes')
parser.add_argument('--dist-url', default='tcp://172.16.1.186:2222', type=str, help='url used to set up distributed training')
parser.add_argument('--dist-backend', default='gloo', type=str, help='distributed backend')
parser.add_argument('--dist-rank', default=0, type=int, help='rank of distributed processes')

这几个是必要的参数设置，其中最后一个是官网没有的

if args.distributed:
    dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url,world_size=args.world_size,rank=args.dist_rank)

这个是分布式的初始化，同样，最后添加一个rank

model.cuda()
model = torch.nn.parallel.DistributedDataParallel(model)

这里，把我们平时使用的单机多卡，数据并行的API

model = torch.nn.DataParallel(model).cuda()

换掉即可。

if args.distributed:
        train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)

最后使用这个官方给的划分方法，把数据集划分即可。

MNIST 分布式例子

MNIST 分布式（2台 GPU机器，每台一张GPU）

GPU1: python mnist_dist.py --init-method=file:///home/workspace/share/1 --rank=0 --world-size=2

GPU2: python mnist_dist.py --init-method=file:///home/workspace/share/1 --rank=1 --world-size=2

from __future__ import print_function
import argparse
import time

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms

import torch.distributed as dist
import torch.utils.data
import torch.utils.data.distributed


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5, 1)
        self.conv2 = nn.Conv2d(20, 50, 5, 1)
        self.fc1 = nn.Linear(4 * 4 * 50, 500)
        self.fc2 = nn.Linear(500, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, 2, 2)
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2, 2)
        x = x.view(-1, 4 * 4 * 50)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return F.log_softmax(x)


def train(args, model, device, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % args.log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                       100. * batch_idx / len(train_loader), loss.item()))


def test(args, model, device, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    for data, target in test_loader:
        data, target = data.to(device), target.to(device)
        output = model(data)
        test_loss += F.nll_loss(output, target, size_average=False).item()  # sum up batch loss
        pred = output.data.max(1, keepdim=True)[1]  # get the index of the max log-probability
        correct += pred.eq(target.data.view_as(pred)).cpu().sum()

    test_loss /= len(test_loader.dataset)

    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))


def main():
    # Training settings
    parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
    parser.add_argument('--batch-size', type=int, default=64, metavar='N',
                        help='input batch size for training (default: 64)')
    parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',
                        help='input batch size for testing (default: 1000)')
    parser.add_argument('--epochs', type=int, default=10, metavar='N',
                        help='number of epochs to train (default: 10)')
    parser.add_argument('--lr', type=float, default=0.01, metavar='LR',
                        help='learning rate (default: 0.01)')
    parser.add_argument('--momentum', type=float, default=0.5, metavar='M',
                        help='SGD momentum (default: 0.5)')
    parser.add_argument('--no-cuda', action='store_true',
                        help='disables CUDA training')
    parser.add_argument('--seed', type=int, default=1, metavar='S',
                        help='random seed (default: 1)')
    parser.add_argument('--log-interval', type=int, default=500, metavar='N',
                        help='how many batches to wait before logging training status')

    parser.add_argument('--init-method', type=str, default='tcp://127.0.0.1:23456')
    parser.add_argument('--rank', type=int)
    parser.add_argument('--world-size', type=int)

    args = parser.parse_args()
    use_cuda = not args.no_cuda and torch.cuda.is_available()
    print(args)

    # 初始化
    dist.init_process_group(init_method=args.init_method, backend="gloo", world_size=args.world_size, rank=args.rank,
                            group_name="pytorch_test")

    torch.manual_seed(args.seed)
    if use_cuda:
        torch.cuda.manual_seed(args.seed)

    train_dataset = datasets.MNIST('./data', train=True, download=False,
                                   transform=transforms.Compose([
                                       transforms.ToTensor(),
                                       transforms.Normalize((0.1307,), (0.3081,))
                                   ]))
    # 分发数据
    train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)

    kwargs = {'num_workers': 5, 'pin_memory': True} if use_cuda else {}

    train_loader = torch.utils.data.DataLoader(train_dataset,
                                               batch_size=args.batch_size, shuffle=True, **kwargs)

    test_loader = torch.utils.data.DataLoader(
        datasets.MNIST('data', train=False, transform=transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))
        ])),
        batch_size=args.test_batch_size, shuffle=True, **kwargs)

    device = torch.device("cuda" if use_cuda else "cpu")
    print(device)
    model = Net().to(device)
    if use_cuda:
        model = torch.nn.parallel.DistributedDataParallel(
            model) if use_cuda else torch.nn.parallel.DistributedDataParallelCPU(model)

    optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)

    total_time = 0
    for epoch in range(1, args.epochs + 1):
        # 设置epoch位置，这应该是个为了同步所做的工作
        train_sampler.set_epoch(epoch)

        start_cpu_secs = time.time()
        train(args, model, device, train_loader, optimizer, epoch)
        end_cpu_secs = time.time()

        print("Epoch {} of {} took {:.3f}s".format(
            epoch, args.epochs, end_cpu_secs - start_cpu_secs))
        total_time += end_cpu_secs - start_cpu_secs

        test(args, model, device, test_loader)

    print("Total time= {:.3f}s".format(total_time))


if __name__ == '__main__':
    main()

单机例子（一台GPU机器一张GPU卡）

python mnist_no_dist.py

from __future__ import print_function
import argparse
import time

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5, 1)
        self.conv2 = nn.Conv2d(20, 50, 5, 1)
        self.fc1 = nn.Linear(4 * 4 * 50, 500)
        self.fc2 = nn.Linear(500, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, 2, 2)
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2, 2)
        x = x.view(-1, 4 * 4 * 50)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)


def train(args, model, device, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % args.log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                       100. * batch_idx / len(train_loader), loss.item()))


def test(args, model, device, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += F.nll_loss(output, target, reduction='sum').item()  # sum up batch loss
            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)

    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))


def main():
    # Training settings
    parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
    parser.add_argument('--batch-size', type=int, default=64, metavar='N',
                        help='input batch size for training (default: 64)')
    parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',
                        help='input batch size for testing (default: 1000)')
    parser.add_argument('--epochs', type=int, default=10, metavar='N',
                        help='number of epochs to train (default: 10)')
    parser.add_argument('--lr', type=float, default=0.01, metavar='LR',
                        help='learning rate (default: 0.01)')
    parser.add_argument('--momentum', type=float, default=0.5, metavar='M',
                        help='SGD momentum (default: 0.5)')
    parser.add_argument('--no-cuda', action='store_true', default=False,
                        help='disables CUDA training')
    parser.add_argument('--seed', type=int, default=1, metavar='S',
                        help='random seed (default: 1)')
    parser.add_argument('--log-interval', type=int, default=500, metavar='N',
                        help='how many batches to wait before logging training status')

    # parser.add_argument('--save-model', action='store_true', default=False,
    #                     help='For Saving the current Model')
    args = parser.parse_args()
    print(args)

    use_cuda = not args.no_cuda and torch.cuda.is_available()

    torch.manual_seed(args.seed)

    kwargs = {'num_workers': 5, 'pin_memory': True} if use_cuda else {}
    train_loader = torch.utils.data.DataLoader(
        datasets.MNIST('./data', train=True, download=True,
                       transform=transforms.Compose([
                           transforms.ToTensor(),
                           transforms.Normalize((0.1307,), (0.3081,))
                       ])),
        batch_size=args.batch_size, shuffle=True, **kwargs)
    test_loader = torch.utils.data.DataLoader(
        datasets.MNIST('./data', train=False, transform=transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))
        ])),
        batch_size=args.test_batch_size, shuffle=True, **kwargs)

    device = torch.device("cuda" if use_cuda else "cpu")
    print(device)
    model = Net().to(device)
    optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)

    total_time = 0
    for epoch in range(1, args.epochs + 1):
        start_cpu_secs = time.time()
        train(args, model, device, train_loader, optimizer, epoch)
        end_cpu_secs = time.time()

        print("Epoch {} of {} took {:.3f}s".format(
            epoch, args.epochs, end_cpu_secs - start_cpu_secs))
        total_time += end_cpu_secs - start_cpu_secs

        test(args, model, device, test_loader)

    # print("Total time= {:.3f}s".format(total_time))
    # if (args.save_model):
    #     torch.save(model.state_dict(), "mnist_cnn.pt")


if __name__ == '__main__':
    main()

参考：

[深度学习] Pytorch 1.0 分布式训练初探