普通单卡训练-GPU

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 distclass 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 outdef train(gpu, args):torch.manual_seed(0)model = ConvNet()torch.cuda.set_device(gpu)model.cuda(gpu)batch_size = 100# define loss function (criterion) and optimizercriterion = nn.CrossEntropyLoss().cuda(gpu)optimizer = torch.optim.SGD(model.parameters(), 1e-4)# Data loading codetrain_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 passoutputs = model(images)loss = criterion(outputs, labels)# Backward and optimizeoptimizer.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))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)if __name__ == '__main__':main()

普通单卡训练-CPU

以上代码在没有GPU的环境上运行会报AssertionError: Torch not compiled with CUDA enabled的错
在程序开始处加上：

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

代码其余地方出现.cuda()的地方改成.to(device)就可以在无gpu的环境中运行了，改完以后就是下面单卡训练——CPU。

CPU单卡训练完整代码如下

# mnist-cpu.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 distdevice = 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):#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)criterion = nn.CrossEntropyLoss()optimizer = torch.optim.SGD(model.parameters(), 1e-4)# Data loading codetrain_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)images = images.to(device)labels = labels.to(device)# Forward passoutputs = model(images)loss = criterion(outputs, labels)# Backward and optimizeoptimizer.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)# 通过启动主函数来开始训练
if __name__ == '__main__':main()

你可能注意到有些参数是多余的，但是对后面的分布式训练是有用的。我们通过执行以下语句就可以在单机单卡上训练：

python mnist-cpu.py -n 1 -g 1 -nr 0

分布式训练-GPU

使用多进程进行分布式训练，我们需要为每个GPU启动一个进程。每个进程需要知道自己运行在哪个GPU上，以及自身在所有进程中的序号。对于多节点，我们需要在每个节点启动脚本。

args.nodes——节点总数
args.gpus——每个节点的GPU总数（每个节点GPU数是一样的）
args.nr ——当前节点在所有节点的序号。
world_size——节点总数乘以每个节点的GPU数可以得到world_size，也即进程总数。
所有的进程需要知道进程0的IP地址以及端口，这样所有进程可以在开始时同步，一般情况下称进程0是master进程，比如我们会在进程0中打印信息或者保存模型。PyTorch提供了mp.spawn来在一个节点启动该节点所有进程，每个进程运行train(i, args)，其中i从0到args.gpus - 1。

# mnist-gpu-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 distdevice = 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 + gpudist.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 optimizercriterion = 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 codetrain_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 passoutputs = model(images)loss = criterion(outputs, labels)# Backward and optimizeoptimizer.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.nodesos.environ['MASTER_ADDR'] = '10.57.23.164' # 服务器ip地址os.environ['MASTER_PORT'] = '8888'mp.spawn(train, nprocs=args.gpus, args=(args,))# 通过启动主函数来开始训练
if __name__ == '__main__':main()

最后就可以执行代码了，比如我们是4节点，每个节点是8卡，那么需要在4个节点分别执行：

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

分布式训练-CPU

# 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 distdevice = 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(cpu, args):rank = args.nr * args.cpus + cpu # 首先计算出当前进程序号：rank = args.nr * args.cpus + cpuprint(rank)dist.init_process_group( # 通过dist.init_process_group初始化分布式环境backend='gloo', # 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)print("111")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)criterion = nn.CrossEntropyLoss()optimizer = torch.optim.SGD(model.parameters(), 1e-4)# Wrap the model# 对于模型侧，我们只需要用DistributedDataParallel包装一下原来的model即可，在背后它会支持梯度的All-Reduce操作。model = nn.parallel.DistributedDataParallel(model)# Data loading codetrain_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)print("333")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)images = images.to(device)labels = labels.to(device)# Forward passoutputs = model(images)loss = criterion(outputs, labels)# Backward and optimizeoptimizer.zero_grad()loss.backward()optimizer.step()if (i + 1) % 100 == 0 and cpu == 0:print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(epoch + 1,args.epochs,i + 1,total_step,loss.item()))if cpu == 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('-c', '--cpus', default=1, type=int,help='number of cpus 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.cpus * args.nodesos.environ['MASTER_ADDR'] = '127.0.0.1'os.environ['MASTER_PORT'] = '8888'mp.spawn(train, nprocs=args.cpus, args=(args,))# 通过启动主函数来开始训练
if __name__ == '__main__':main()

分布式训练方式：

使用两台装有centos7系统的虚拟机，确保防火墙已关闭（防火墙开启会导致无法同时训练）
关闭防火墙：systemctl stop firewalld
检查防火墙状态：systemctl status firewalld
保证两台虚拟机能互相免密登录（即ssh+主机名，不需要输入密码就能登录另一台机器），免密登录教程：linux虚拟机之间ssh免密登录配置
将os.environ['MASTER_ADDR'] = '127.0.0.1'中的master ip改为真正的master ip，可以通过ifconfig查看
分别通过以下命令开始训练
python mnist-cpu-distributed.py -n 2 -nr 0
python mnist-cpu-distributed.py -n 2 -nr 1

训练结果如下图，原本在一台机器训练时每次迭代需要600 steps，现在分布式训练每台虚拟机只需要300 steps
如果报错address is already in use，说明已经有进程占用端口，重启虚拟机或杀死进程解决（重启虚拟机之后再次检查防火墙是否关闭）
```
netstat -tunlp | grep 端口 #根据端口查看占用进程
kill -9 进程号 #杀掉进程
```

参考链接：Pytorch分布式训练简明教程

租GPU服务器相关

本机没有GPU环境的小伙伴可以租用MistGPU服务器进行训练，教程：【Pytorch分布式训练】MistGPU服务器训练

【Pytorch分布式训练】在MNIST数据集上训练一个简单CNN网络，将其改成分布式训练相关推荐

在MNIST数据集上训练一个手写数字识别模型
使用Pytorch在MNIST数据集上训练一个手写数字识别模型, 代码和参数文件可下载 1.1 数据下载 import torchvision as tvtraining_sets = tv.dat ...
DL之DCGNN：基于TF利用DCGAN实现在MNIST数据集上训练生成新样本
DL之DCGNN:基于TF利用DCGAN实现在MNIST数据集上训练生成新样本目录输出结果设计思路实现部分代码说明:所有图片文件丢失输出结果更新-- 设计思路更新-- 实现部分代码更 ...
tensorflow（七）实现mnist数据集上图片的训练和测试
本文使用tensorflow实现在mnist数据集上的图片训练和测试过程,使用了简单的两层神经网络,代码中涉及到的内容,均以备注的形式标出. 关于文中的数据集,大家如果没有下载下来,可以到我的网盘去下 ...
使用mnist数据集_使用MNIST数据集上的t分布随机邻居嵌入（t-SNE）进行降维
使用mnist数据集 It is easy for us to visualize two or three dimensional data, but once it goes beyond thr ...
autoencoder自编码器原理以及在mnist数据集上的实现
Autoencoder是常见的一种非监督学习的神经网络.它实际由一组相对应的神经网络组成(可以是普通的全连接层,或者是卷积层,亦或者是LSTMRNN等等,取决于项目目的),其目的是将输入数据降维成一个 ...
在一行上添加一个简单的if-then-else语句[重复]
本文翻译自:Putting a simple if-then-else statement on one line [duplicate] Possible Duplicate: 可能重复: Pyth ...
在Kubernetes上部署一个简单的、类PaaS的平台，原来这么容易！
作者 | Bram Dingelstad 译者 | 弯月责编 |徐威龙封图| CSDN下载于视觉中国我们都遇到过这种情况:有人发现了一个bug,然而这不是一般的软件bug,甚至都不是通常意义上的 ...
keras笔记-mnist数据集上的简单训练
学习了keras已经好几天了,之前一直拒绝使用keras,但是现在感觉keras是真的好用啊,可以去尝试一下啊. 首先展示的第一个代码,还是mnist数据集的训练和测试,以下是代码: from ker ...
深度学习笔记（2）——pytorch实现MNIST数据集分类（FNN、CNN、RNN、LSTM、GRU）
文章目录 0 前言 1 数据预处理 2 FNN(前馈神经网络) 3 CNN(卷积神经网络) 4 RNN(循环神经网络) 5 LSTM(长短期记忆网络) 6 GRU(门控循环单元) 7 完整代码 0 前 ...

【Pytorch分布式训练】在MNIST数据集上训练一个简单CNN网络，将其改成分布式训练

文章目录