加强红色网站建设,高端网站开发设计,厦门旅游网站设计,建筑网站源码Pytorch单机多卡分布式训练
数据并行#xff1a;
DP和DDP
这两个都是pytorch下实现多GPU训练的库#xff0c;DP是pytorch以前实现的库#xff0c;现在官方更推荐使用DDP#xff0c;即使是单机训练也比DP快。 DataParallel#xff08;DP#xff09; 只支持单进程多线程…Pytorch单机多卡分布式训练
数据并行
DP和DDP
这两个都是pytorch下实现多GPU训练的库DP是pytorch以前实现的库现在官方更推荐使用DDP即使是单机训练也比DP快。 DataParallelDP 只支持单进程多线程单一机器上进行训练。模型训练开始的时候先把模型复制到四个GPU上面然后把数据分配给四个GPU进行前向传播前向传播之后再汇总到卡0上面然后在卡0上进行反向传播参数更新再将更新好的模型复制到其他几张卡上。 DistributedDataParallelDDP 支持多线程多进程单一或者多个机器上进行训练。通常DDP比DP要快。 先把模型载入到四张卡上每个GPU上都分配一些小批量的数据再进行前向传播反向传播计算完梯度之后再把所有卡上的梯度汇聚到卡0上面卡0算完梯度的平均值之后广播给所有的卡所有的卡更新自己的模型这样传输的数据量会少很多。
DDP代码写法 初始化 import torch.distributed as dist
import torch.utils.data.distributed# 进行初始化backend表示通信方式可选择的有nccl英伟达的GPU2GPU的通信库适用于具有英伟达GPU的分布式训练、gloo基于tcp/ip的后端可在不同机器之间进行通信通常适用于不具备英伟达GPU的环境、mpi适用于支持mpi集群的环境
# init_method: 告知每个进程如何发现彼此默认使用env://
dist.init_process_group(backendnccl, init_methodenv://)设置device device torch.device(fcuda:{args.local_rank}) # 设置devicelocal_rank表示当前机器的进程号该方式为每个显卡一个进程
torch.cuda.set_device(device) # 设定device创建dataloader之前要加一个sampler trans transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
data_set torchvision.datasets.MNIST(./, trainTrue, transformtrans, target_transformNone, downloadTrue)
train_sampler torch.utils.data.distributed.DistributedSampler(data_set) # 加一个sampler
data_loader_train torch.utils.data.DataLoader(datasetdata_set, batch_size256, samplertrain_sampler)torch.nn.parallel.DistributedDataParallel包裹模型先to(device)再包裹模型 net torchvision.models.resnet101(num_classes10)
net.conv1 torch.nn.Conv2d(1, 64, (7, 7), (2, 2), (3, 3), biasFalse)
net net.to(device)
net torch.nn.parallel.DistributedDataParallel(net, device_ids[device], output_device[device]) # 包裹模型真正训练之前要set_epoch()否则将不会shuffer数据 for epoch in range(10):train_sampler.set_epoch(epoch) # set_epochfor step, data in enumerate(data_loader_train):images, labels dataimages, labels images.to(device), labels.to(device)opt.zero_grad()outputs net(images)loss criterion(outputs, labels)loss.backward()opt.step()if step % 10 0:print(loss: {}.format(loss.item()))模型保存 if args.local_rank 0: # local_rank为0表示master进程torch.save(net, my_net.pth)运行 if __name__ __main__:parser argparse.ArgumentParser()# local_rank参数是必须的运行的时候不必自己指定DDP会自行提供parser.add_argument(--local_rank, typeint, default0)args parser.parse_args()main(args)运行命令 python -m torch.distributed.launch --nproc_per_node2 多卡训练.py # --nproc_per_node2表示当前机器上有两个GPU可以使用完整代码
import os
import argparse
import torch
import torchvision
import torch.distributed as dist
import torch.utils.data.distributedfrom torchvision import transforms
from torch.multiprocessing import Processdef main(args):# nccl: 后端基于NVIDIA的GPU-to-GPU通信库适用于具有NVIDIA GPU的分布式训练# gloo: 后端是一个基于TCP/IP的后端可在不同机器之间进行通信通常适用于不具备NVIDIA GPU的环境。# mpi 后端使用MPI实现适用于具备MPI支持的集群环境。# init_method: 告知每个进程如何发现彼此如何使用通信后端初始化和验证进程组。 默认情况下如果未指定 init_methodPyTorch 将使用环境变量初始化方法 (env://)。dist.init_process_group(backendnccl, init_methodenv://) # nccl比较推荐device torch.device(fcuda:{args.local_rank})torch.cuda.set_device(device)trans transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])data_set torchvision.datasets.MNIST(./, trainTrue, transformtrans, target_transformNone, downloadTrue)train_sampler torch.utils.data.distributed.DistributedSampler(data_set)data_loader_train torch.utils.data.DataLoader(datasetdata_set, batch_size256, samplertrain_sampler)net torchvision.models.resnet101(num_classes10)net.conv1 torch.nn.Conv2d(1, 64, (7, 7), (2, 2), (3, 3), biasFalse)net net.to(device)net torch.nn.parallel.DistributedDataParallel(net, device_ids[device], output_device[device])criterion torch.nn.CrossEntropyLoss()opt torch.optim.Adam(paramsnet.parameters(), lr0.001)for epoch in range(10):train_sampler.set_epoch(epoch)for step, data in enumerate(data_loader_train):images, labels dataimages, labels images.to(device), labels.to(device)opt.zero_grad()outputs net(images)loss criterion(outputs, labels)loss.backward()opt.step()if step % 10 0:print(loss: {}.format(loss.item()))if args.local_rank 0:torch.save(net, my_net.pth)if __name__ __main__:parser argparse.ArgumentParser()# must parse the command-line argument: --local_rankLOCAL_PROCESS_RANK, which will be provided by DDPparser.add_argument(--local_rank, typeint, default0)args parser.parse_args()main(args)
参考
https://zhuanlan.zhihu.com/p/594046884
https://zhuanlan.zhihu.com/p/358974461
