PyTorch常用函数(7)

1. 可视化层

• torch.nn.PixelShuffle(upscale_factor)：
  将 shape 为$[N,C\ast r^2,H,W]$的 Tensor 重新排列为 shape 为$[N,C,H\ast r,W\ast r]$的 Tensor。 当使用stride=1/r 的 sub-pixel 卷积的时候，这个方法是非常有用的。
  参数
  （1）upscale_factor (int)：增加空间分辨率因子
  形状
  （1）Input：$(\ast ,C_{in},H_{in},W_{in})$
  （2）Output：$(\ast ,C_{out},H_{out},W_{out})$
  $$\begin{gathered} C_{out}=C_{in}/upscale\_facto{r}^2 \\ {H}_{out}=H_{in}\ast upscale\_factor \\ {W}_{out}=W_{in}\ast upscale\_factor \end{gathered}$$

import torch
import torch.nn as nn
pixel_shuffle = nn.PixelShuffle(3)
input = torch.randn(1, 9, 4, 4)
output = pixel_shuffle(input)
print(output.size())
# torch.Size([1, 1, 12, 12])

• torch.nn.UpsamplingNearest2d(size=None, scale_factor=None)：对于多 channel 输入 进行 2-D 最近邻上采样。可以通过 size 或者 scale_factor 来指定上采样后的图片大小。当给定 size 时， size 的值将会是输出图片的大小。
  参数：
  （1）size (tuple, optional) – 一个包含两个整数的元组 $(H_{out},W_{out})$指定了输出的长宽
  （2）scale_factor (int, optional) – 长和宽的一个乘子
  形状
  （1）Input:$(N,C,H_{in},W_{in})$
  （2）Output: $(N,C,H_{out},W_{out})$
  $H_{out}=floor(H_{in}\ast scale\_factor)$
  $W_{out}=floor(W_{in}\ast scale\_factor)$

import torch
from torch import nn
input = torch.arange(1, 5, dtype=torch.float32).view(1, 1, 2, 2)
print(input)
# tensor([[[[1., 2.],
#           [3., 4.]]]])
m = nn.UpsamplingNearest2d(scale_factor=2)
print(m(input))
# tensor([[[[ 1.,  1.,  2.,  2.],
#           [ 1.,  1.,  2.,  2.],
#           [ 3.,  3.,  4.,  4.],
#           [ 3.,  3.,  4.,  4.]]]])

• torch.nn.UpsamplingBilinear2d(size=None, scale_factor=None)：对于多 channel 输入 进行 2-D bilinear 上采样。
  可以通过 size 或者 scale_factor 来指定上采样后的图片大小。
  当给定 size 时， size 的值将会是输出图片的大小。
  参数：
  （1）size (tuple, optional)： 一个包含两个整数的元组 $(H_{out},W_{out})$指定了输出的长宽
  （2）scale_factor (int, optional)：长和宽的一个乘子
  形状
  （1）Input:$(N,C,H_{in},W_{in})$
  （2）Output: $(N,C,H_{out},W_{out})$
  $H_{out}=floor(H_{in}\ast scale\_factor)$
  $W_{out}=floor(W_{in}\ast scale\_factor)$

import torch
from torch import nn
input = torch.arange(1, 5, dtype=torch.float32).view(1, 1, 2, 2)
print(input)
# tensor([[[[1., 2.],
#           [3., 4.]]]])
m = nn.UpsamplingBilinear2d(scale_factor=2)
print(m(input))
# tensor([[[[1.0000, 1.3333, 1.6667, 2.0000],
#           [1.6667, 2.0000, 2.3333, 2.6667],
#           [2.3333, 2.6667, 3.0000, 3.3333],
#           [3.0000, 3.3333, 3.6667, 4.0000]]]])

2. 多GPU层

• torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0)：在模块级别上实现数据并行。
  通过将 mini-batch 划分到不同的设备上来实现给定 module 的并行。在 forward 过程中， module会在每个设备上都复制一遍，每个副本都会处理部分输入。在反向传播时，每个卡上的梯度会汇总到主卡上,求得梯度的均值后,再用反向传播更新单个GPU上的模型参数，最后将更新后的模型参数复制到剩余指定的GPU中进行下一轮的前向传播,以此来实现并行。
  batch 的大小应该大于所使用的 GPU 的数量。还应当是 GPU 个数的整数倍，这样划分出来的每一块都会有相同的样本数量。
  参数说明：
  （1）module – 要放到多卡训练的模型
  （2）device_ids –数据类型是一个列表, 表示可用的gpu卡号，默认为所有设备。
  （3）output_device – 输出设备（默认为 device_ids[0]），这也是为什么多gpu训练并不是负载均衡的,一般0卡会占用的多,这里还涉及到一个小知识点——如果程序开始加os.environ[“CUDA_VISIBLE_DEVICES”] = “2, 3”, 那么0卡(逻辑卡号)指的是2卡(物理卡号))。

import torch
import torch.nn as nn
os.environ['CUDA_VISIBLE_DEVICES']='2, 1' # 指定逻辑卡号。比如这样写的话,第0块卡指的是物理上的第2块

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
input = torch.randn(8, 4, 160, 160).to(device) # 将数据放到gpu上
model = nn.Conv2d(4, 3, kernel_size = 3, stride = 1, padding = 1)
model = model.to(device)# 将模型放到gpu上
output = model(input) 
model = torch.nn.DataParallel(model, device_ids=[0, 1])

• torch.nn.parallel.DistributedDataParallel(module, device_ids=None, output_device=None, dim=0, broadcast_buffers=True, process_group=None, bucket_cap_mb=25, find_unused_parameters=False, check_reduction=False)：

参数定义
（1）module是要放到多卡训练的模型;
（2）device_ids数据类型是一个列表, 表示可用的gpu卡号;
（3）output_devices数据类型也是列表,表示模型输出结果存放的卡号(如果不指定的话,默认放在0卡,这也是为什么多gpu训练并不是负载均衡的,一般0卡会占用的多,这里还涉及到一个小知识点——如果程序开始加os.environ[“CUDA_VISIBLE_DEVICES”] = “2, 3”, 那么0卡(逻辑卡号)指的是2卡(物理卡号))。
（4）dim指按哪个维度进行数据的划分,默认是输入数据的第一个维度,即按batchsize划分(设数据数据的格式是B, C, H, W)

分布式的几个概念
（1）group:即进程组。默认情况下，只有一个组，一个 job 即为一个组，也即一个 world。 当需要进行更加精细的通信时，可以通过 new_group 接口，使用 word 的子集，创建新组，用于集体通信等;
（2）world size:表示全局进程个数;
（3）rank:表示进程序号，用于进程间通讯，表征进程优先级。rank = 0 的GPU为主卡;
（4）local_rank:是当前主机节点当前的进程号，而非进程内的gpu号，在脚本启动后会自动设置
DDP的原理
DDP在各进程梯度计算完成之后,各进程需要将梯度进行汇总平均,然后再由 rank=0 的进程,将其 broadcast 到所有进程后,各进程用该梯度来独立的更新参数而 DP是梯度汇总到GPU0,反向传播更新参数,再广播参数给其他剩余的GPU。由于DDP各进程中的模型,初始参数一致 (初始时刻进行一次 broadcast),而每次用于更新参数的梯度也一致,因此,各进程的模型参数始终保持一致。而在DP中,全程维护一个 optimizer,对各个GPU上梯度进行求平均,而在主卡进行参数更新,之后再将模型参数 broadcast 到其他GPU.相较于DP, DDP传输的数据量更少,因此速度更快,效率更高。
代码解释

# 引入包
import argparse
import torch.distributed as dist

# 设置可选参数
parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', default=0, type=int,
                    help='node rank for distributed training')
args = parser.parse_args()
# print(args.local_rank)
dist.init_process_group(backend='nccl')


# 1.初始化进程组
dist.init_process_group(backend='nccl')
torch.cuda.set_device(args.local_rank)

# 2.使用DistributedSampler
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle = (train_sampler is None), sampler=train_sampler, pin_memory=False)

# 3.创建DDP模型进行分布式训练
model = nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank, find_unused_parameters=True)

# 4.命令行开始训练 --nproc_per_node参数指定为当前主机创建的进程数(比如我当前可用但卡数是2 那就为这个主机创建两个进程,每个进程独立执行训练脚本)
# 我是单机多卡, 所以nnode=1, 就是一台主机, 一台主机上--nproc_per_node个进程
python -m torch.distributed.launch --nnodes=1 --nproc_per_node=2 --node_rank=0  --master_port=6005 train.py

1）init_process_group初始化进程组(如果需要进行前面提到的小组内集体通信,用new_group创建子分组);

# 初始化使用nccl后端,官网还提供了‘gloo’作为backend,nccl是用来GPU之间进行通信的
torch.distributed.init_process_group(backend="nccl")

2）使用DistributedSampler(改dataloader中sampler=train_sampler)

# train_dataset就是自己的数据集
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle = (train_sampler is None), sampler=train_sampler, pin_memory=False)

pin_memory(固定内存)设为True会加速训练
3）创建DDP模型进行分布式训练

# DDP是all-reduce的,即汇总不同 GPU 计算所得的梯度,并同步计算结果。all-reduce 后不同 GPU 中模型的梯度均为 all-reduce 之前各 GPU 梯度的均值.
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank, find_unused_parameters=True)

4）最后一步,执行python train.py时需加参数

python -m torch.distributed.launch --nnodes=1 --nproc_per_node=2 --node_rank=0  --master_port=6005 train.py

• nn.DataParallel(以下简称DP)和DistributedDataParallel(以下简称DDP)的区别

1）DDP通过多进程实现的。也就是说操作系统会为每个GPU创建一个进程,从而避免了Python解释器GIL带来的性能开销。而DataParallel()是通过单进程控制多线程来实现的。还有一点,DDP也不存在前面DP提到的负载不均衡问题。

2）参数更新的方式不同。DDP在各进程梯度计算完成之后,各进程需要将梯度进行汇总平均,然后再由 rank=0 的进程,将其 broadcast 到所有进程后,各进程用该梯度来独立的更新参数而 DP是梯度汇总到GPU0,反向传播更新参数,再广播参数给其他剩余的GPU。也就是说，在每个训练批次（batch）中，因为模型的权重都是在 一个进程上先算出来 然后再把他们分发到每个GPU上，所以网络通信就成为了一个瓶颈，而GPU使用率也通常很低。由于DDP各进程中的模型,初始参数一致 (初始时刻进行一次 broadcast),而每次用于更新参数的梯度也一致,因此,各进程的模型参数始终保持一致。而在DP中,全程维护一个 optimizer,对各个GPU上梯度进行求平均,而在主卡进行参数更新,之后再将模型参数 broadcast 到其他GPU。DP是传输更新后的参数，DDP是传输梯度，但是实际上做的时候并不是直接传输梯度，是利用RingAllReduce算法做了加速求平均，传输量小很多，因此速度更快,效率更高。

3）DDP支持 all-reduce(指汇总不同 GPU 计算所得的梯度,并同步计算结果),broadcast,send 和 receive 等等。通过 MPI 实现 CPU 通信,通过 NCCL 实现 GPU 通信,缓解了“写在前面“提到的进程间通信有大的开销问题。
在DataParallel中,batch_size设置必须为单卡的n倍, 但是在DistributedDataParallel内,batch_size设置于单卡一样即可。实验发现 单GPU跑但是套DP的话 指标会降低很多。

参考目录

https://blog.csdn.net/weixin_40920183/article/details/119814472
https://zhuanlan.zhihu.com/p/206467852