《深入浅出PyTorch》学习笔记——第二章：PyTorch基础知识

从本章开始，我们将开始介绍PyTorch基础知识，本章我们将介绍张量，以帮助大家建立起对数据的描述，随后我们再介绍张量的运算，最后再讲PyTorch中所有神经网络的核心包 autograd ，也就是自动微分，了解完这些内容我们就可以较好地理解PyTorch代码了。在深度学习中，我们通常将数据以张量的形式进行表示，比如我们用三维张量表示一个RGB图像，四维张量表示视频。

经过本节的学习，你将收获：

张量的简介
PyTorch如何创建张量
PyTorch中张量的操作
PyTorch中张量的广播机制

1.1 简介

几何代数中定义的张量是基于向量和矩阵的推广，比如我们可以将标量视为零阶张量，矢量可以视为一阶张量，矩阵就是二阶张量。

张量维度	代表含义
0维张量	代表的是标量（数字）
1维张量	代表的是向量
2维张量	代表的是矩阵
3维张量	时间序列数据股价文本数据单张彩色图片(RGB)

张量是现代机器学习的基础。它的核心是一个数据容器，多数情况下，它包含数字，有时候它也包含字符串，但这种情况比较少。因此可以把它想象成一个数字的水桶。

这里有一些存储在各种类型张量的公用数据集类型：

3维 = 时间序列
4维 = 图像
5维 = 视频

例子：一个图像可以用三个字段表示：

(width, height, channel) = 3D

但是，在机器学习工作中，我们经常要处理不止一张图片或一篇文档——我们要处理一个集合。我们可能有10,000张郁金香的图片，这意味着，我们将用到4D张量：

(batch_size, width, height, channel) = 4D

在PyTorch中， torch.Tensor 是存储和变换数据的主要工具。如果你之前用过NumPy，你会发现 Tensor 和NumPy的多维数组非常类似。然而，Tensor 提供GPU计算和自动求梯度等更多功能，这些使 Tensor 这一数据类型更加适合深度学习。

1.2 tensor和numpy中ndarray的区别

张量（Tensor） 和 NumPy 中的多维数组（ndarray） 是多维数据的两种表示形式，它们在功能、使用场景和性能优化方面有一些区别。以下是二者的主要差异：

1.3 tensor的常用函数

其实torch的学习就是不断查阅官方文档巩固的过程，官网对于tensor的介绍非常详细，但很多函数我们平时基本都用不上，真到需要用的时候去查阅即可。

官方文档链接：torch — PyTorch 2.5 documentation

以下仅举出一些常用的函数：

1. 创建相关

zeros: 创建一个全为零的 Tensor。
ones: 创建一个全为一的 Tensor。
arange: 创建一个从起始值到终止值按固定步长的 1-D Tensor。
linspace: 创建一个在指定范围内均匀分布的 1-D Tensor。
empty: 创建一个未初始化的 Tensor。
as_tensor: 将数据转换为 Tensor，尽可能共享数据和保留 Autograd 历史。
from_numpy: 从 numpy 数组创建 Tensor，共享内存。
rand：随机创建一个指定形状的Tensor，数据范围为[0, 1)。
randn：随机创建一个指定形状的Tensor，数据符合正态分布。
tensor：通过复制数据构造一个没有自grad历史的张量。

2. 索引与切片

where: 根据条件从两个 Tensor 中选择元素。
nonzero: 返回输入 Tensor 中非零元素的索引。
index_select: 根据索引在指定维度上选择 Tensor 的元素。
masked_select: 根据布尔掩码选择 Tensor 的元素。

3. 拼接与分割

cat: 沿指定维度拼接一组 Tensor。
stack: 沿新维度拼接一组 Tensor。
split: 将 Tensor 拆分为若干子 Tensor。
chunk: 将 Tensor 分块为指定数量的小块。

4. 形状调整

reshape: 改变 Tensor 的形状而不改变数据。
permute: 根据指定维度顺序重新排列 Tensor。
transpose: 转置 Tensor 的两个维度。
squeeze: 删除指定的单维度。
unsqueeze: 在指定位置插入单维度。
view：view()仅仅是改变了对这个张量的观察角度，两者共享内存。

5. 元素统计与属性

is_tensor: 判断对象是否为 PyTorch Tensor。
is_floating_point: 判断输入 Tensor 数据类型是否为浮点型。
numel: 返回 Tensor 中的总元素数量。

6. 数值操作

add: 对 Tensor 进行加法运算。
sub: 对 Tensor 进行减法运算。
mul: 对 Tensor 进行乘法运算。
div: 对 Tensor 进行除法运算。

7. 广播与其他操作

tile: 重复 Tensor 的元素以创建新 Tensor。
expand: 将 Tensor 扩展到指定形状而不复制数据。

8. 转换

zeros_like: 将现有矩阵转换为全0矩阵。
new_ones：根据现有矩阵转换为一个同样形状的全1矩阵。
randn_like：返回一个与输入大小相同的张量，其中填充了来自均值为0、方差为1的正态分布的随机数。

1.4 创建tensor

在接下来的内容中，我们将介绍几种常见的创建tensor的方法。

1）随机初始化矩阵

我们可以通过torch.rand()的方法，构造一个随机初始化的矩阵：

import torch
x = torch.rand(2,3)
print(x)

tensor([[0.2857, 0.3249, 0.1768],
[0.3912, 0.2394, 0.5674]])

2）全0矩阵的构建

我们可以通过torch.zeros()构造一个矩阵全为 0，并且通过dtype设置数据类型为 long。除此以外，我们还可以通过torch.zero_()和torch.zeros_like()将现有矩阵转换为全0矩阵.

y = torch.zeros(2, 3, dtype=torch.long)
print(y)

x = torch.zeros_like(x)
print(x)

tensor([[0, 0, 0],
[0, 0, 0]])
tensor([[0., 0., 0.],
[0., 0., 0.]])

3）张量的构建

我们可以通过torch.tensor()直接使用数据，构造一个张量：

z = torch.tensor([2, 3, 4])
print(z)

tensor([2, 3, 4])

4）基于已经存在的 tensor，创建一个 tensor

y = y.new_ones(2, 3, dtype=torch.double)
print(y)

y = torch.randn_like(y, dtype=torch.float)
print(y)

tensor([[1., 1., 1.],
[1., 1., 1.]], dtype=torch.float64)
tensor([[ 0.8883, 2.3753, -2.0295],
[-0.5746, 1.4508, 0.0397]])

5）常见的构造Tensor的方法

函数	功能
Tensor(sizes)	基础构造函数
tensor(data)	类似于np.array
ones(sizes)	全1
zeros(sizes)	全0
eye(sizes)	对角为1，其余为0
arange(s,e,step)	从s到e，步长为step
linspace(s,e,steps)	从s到e，均匀分成step份
rand/randn(sizes)	rand是[0,1)均匀分布；randn是服从N(0，1)的正态分布
normal(mean,std)	正态分布(均值为mean，标准差是std)
randperm(m)	随机排列

1.5 张量的操作

在接下来的内容中，我们将介绍几种常见的张量的操作方法：

1）加法操作：

x = torch.rand(2, 3)
y = torch.rand(2, 3)
print(x)
print(y)
# 方式1:
print(torch.add(x, y))
# 方式2:
y.add_(x)
print(y)

tensor([[0.0515, 0.3962, 0.9322],
[0.2538, 0.6927, 0.3299]])
tensor([[0.5487, 0.1732, 0.3093],
[0.3079, 0.2375, 0.9740]])
tensor([[0.6002, 0.5694, 1.2416],
[0.5617, 0.9302, 1.3038]])
tensor([[0.6002, 0.5694, 1.2416],
[0.5617, 0.9302, 1.3038]])

2）索引操作：(类似于numpy)

需要注意的是：索引出来的结果与原数据共享内存，修改一个，另一个会跟着修改。如果不想修改，可以考虑使用copy()等方法

x = torch.rand(2, 3)
print(x)

y = x[:, 1]
print(y)

y+=1
print(y)
print(x)

tensor([[0.3335, 0.6017, 0.6432],
[0.1937, 0.5229, 0.1899]])
tensor([0.6017, 0.5229])
tensor([1.6017, 1.5229])
tensor([[0.3335, 1.6017, 0.6432],
[0.1937, 1.5229, 0.1899]])

3）维度变换

张量的维度变换常见的方法有torch.view()和torch.reshape()，下面我们将介绍第一中方法torch.view()：

x = torch.randn(4,4)
print(x.shape)
y = x.view(16)
print(y.shape)
z = x.view(-1,8)
print(z.shape)

torch.Size([4, 4])
torch.Size([16])
torch.Size([2, 8])

注: torch.view() 返回的新tensor与源tensor共享内存(其实是同一个tensor)，更改其中的一个，另外一个也会跟着改变。(顾名思义，view()仅仅是改变了对这个张量的观察角度)

x+=1
print(x)
print(y)

tensor([[ 1.3717, 2.1981, 1.0409, 0.8880],
[ 0.2683, 2.4257, -0.7726, 1.4205],
[ 1.5327, 1.0613, 0.7309, 0.3195],
[ 1.1992, 1.8980, 1.2854, 0.8150]])
tensor([ 1.3717, 2.1981, 1.0409, 0.8880, 0.2683, 2.4257, -0.7726, 1.4205,
1.5327, 1.0613, 0.7309, 0.3195, 1.1992, 1.8980, 1.2854, 0.8150])

上面我们说过torch.view()会改变原始张量，但是很多情况下，我们希望原始张量和变换后的张量互相不影响。为为了使创建的张量和原始张量不共享内存，我们需要使用第二种方法torch.reshape()，同样可以改变张量的形状，但是此函数并不能保证返回的是其拷贝值，所以官方不推荐使用。推荐的方法是我们先用 clone() 创造一个张量副本然后再使用 torch.view()进行函数维度变换。

注：使用 clone() 还有一个好处是会被记录在计算图中，即梯度回传到副本时也会传到源 Tensor 。

x = torch.randn(1)
print(type(x))
print(type(x.item()))

<class 'torch.Tensor'>
<class 'float'>

PyTorch中的 Tensor 支持超过一百种操作，包括转置、索引、切片、数学运算、线性代数、随机数等等，具体使用方法可参考官方文档。

1.6 广播机制

当对两个形状不同的 Tensor 按元素运算时，可能会触发广播(broadcasting)机制：先适当复制元素使这两个 Tensor 形状相同后再按元素运算。

x = torch.arange(1, 3).view(1, 2)
print(x)
y = torch.arange(1, 4).view(3, 1)
print(y)
print(x + y)

tensor([[1, 2]])
tensor([[1],
[2],
[3]])
tensor([[2, 3],
[3, 4],
[4, 5]])

由于x和y分别是1行2列和3行1列的矩阵，如果要计算x+y，那么x中第一行的2个元素被广播 (复制)到了第二行和第三行，⽽y中第⼀列的3个元素被广播(复制)到了第二列。如此，就可以对2个3行2列的矩阵按元素相加。

1.7 tensor和numpy中ndarray的相互转换

在1.2我们介绍了“tensor和numpy中ndarray的区别”的区别，那么两者之间该如何转换呢？

1）由张量变换为Numpy array数组

t = torch.ones(5)
print(type(t))
n = t.numpy()
print(type(n))

<class 'torch.Tensor'>
<class 'numpy.ndarray'>

修改张量的值，则Numpy array数组值也会随之改变。

t.add_(1)
print(f"t: {t}")
print(f"n: {n}")

t: tensor([2., 2., 2., 2., 2.])
n: [2. 2. 2. 2. 2.]

2）由Numpy array数组变换为张量

import numpy as np
n = np.ones(5)
print(type(n))
t = torch.from_numpy(n)
print(type(t))

<class 'numpy.ndarray'>
<class 'torch.Tensor'>

修改Numpy array数组的值，则张量值也会随之改变。

np.add(n, 1, out=n)
print(f"t: {t}")
print(f"n: {n}")

t: tensor([2., 2., 2., 2., 2.], dtype=torch.float64)
n: [2. 2. 2. 2. 2.]

2 自动求导简介

PyTorch 中，所有神经网络的核心是 autograd 包。autograd包为张量上的所有操作提供了自动求导机制。它是一个在运行时定义 ( define-by-run ）的框架，这意味着反向传播是根据代码如何运行来决定的，并且每次迭代可以是不同的。

经过本节的学习，你将收获：

autograd的求导机制
梯度的反向传播

2.1 Autograd简介

torch.Tensor 是这个包的核心类。如果设置它的属性.requires_grad 为 True，那么它将会追踪对于该张量的所有操作。当完成计算后可以通过调用.backward()，来自动计算所有的梯度。这个张量的所有梯度将会自动累加到.grad属性。

注意：在 y.backward() 时，如果 y 是标量，则不需要为 backward() 传入任何参数；否则，需要传入一个与 y 同形的Tensor。

要阻止一个张量被跟踪历史，可以调用.detach()方法将其与计算历史分离，并阻止它未来的计算记录被跟踪。为了防止跟踪历史记录(和使用内存），可以将代码块包装在 with torch.no_grad(): 中。在评估模型时特别有用，因为模型可能具有 requires_grad = True 的可训练的参数，但是我们不需要在此过程中对他们进行梯度计算。

还有一个类对于autograd的实现非常重要：Function。Tensor 和Function 互相连接生成了一个无环图 (acyclic graph)，它编码了完整的计算历史。每个张量都有一个.grad_fn属性，该属性引用了创建 Tensor 自身的Function(除非这个张量是用户手动创建的，即这个张量的grad_fn是 None )。下面给出的例子中，张量由用户手动创建，因此grad_fn返回结果是None。

import torch

x = torch.randn(3, 3, requires_grad=True)
print(x.grad_fn)

None

如果需要计算导数，可以在 Tensor 上调用 .backward()。如果Tensor 是一个标量(即它包含一个元素的数据），则不需要为 backward() 指定任何参数，但是如果它有更多的元素，则需要指定一个gradient参数，该参数是形状匹配的张量。

创建一个张量并设置requires_grad=True用来追踪其计算历史

x = torch.ones(2, 2, requires_grad=True)
print(x)

y = x**2
print(y)
print(y.grad_fn)

z = y * y * 3
out = z.mean()

print(z, out)

tensor([[1., 1.],
[1., 1.]], requires_grad=True)
tensor([[1., 1.],
[1., 1.]], grad_fn=<PowBackward0>)
<PowBackward0 object at 0x00000201B0B4EAC0>
tensor([[3., 3.],
[3., 3.]], grad_fn=<MulBackward0>) tensor(3., grad_fn=<MeanBackward0>)

a = torch.randn(2, 2)  # 缺失情况下默认 requires_grad = False
a = (a * 3) / (a - 1)
print(a.requires_grad)
a.requires_grad_(True)
print(a.requires_grad)
b = (a * a).sum()
print(b.grad_fn)

.requires_grad_(...) 原地改变了现有张量的requires_grad标志。如果没有指定的话，默认输入的这个标志是False。

False
True
<SumBackward0 object at 0x00000201B1FFB670>

2.2 梯度

现在开始进行反向传播，因为out 是一个标量，因此out.backward()和out.backward(torch.tensor(1.)) 等价。

输出导数d(out)/dx

out.backward()
print(x.grad)

tensor([[3., 3.],
[3., 3.]])

也即autograd这个包只能用于计算雅可比矩阵的乘积，不能给出完整的雅可比矩阵。

注意：grad在反向传播过程中是累加的(accumulated)，这意味着每一次运行反向传播，梯度都会累加之前的梯度，所以一般在反向传播之前需把梯度清零。

# 再来反向传播⼀一次，注意grad是累加的
out2 = x.sum()
out2.backward()
print(x.grad)

out3 = x.sum()
x.grad.data.zero_()
out3.backward()
print(x.grad)

tensor([[4., 4.],
[4., 4.]])
tensor([[1., 1.],
[1., 1.]])

现在我们来看一个雅可比向量积的例子：

x = torch.ones(3, requires_grad=True)
print(x)
y = 2 * x
# y = x.sum()
v = torch.tensor([0.1, 1.0, 0.0001], dtype=torch.float)
y.backward(v)
print(x.grad)

在这种情况下，y 不再是标量。torch.autograd 不能直接计算完整的雅可比矩阵，但是如果我们只想要雅可比向量积，只需将这个向量作为参数传给 backward：

tensor([1., 1., 1.], requires_grad=True)
tensor([2.0000e-01, 2.0000e+00, 2.0000e-04])

也可以通过将代码块包装在with torch.no_grad(): 中，来阻止 autograd 跟踪设置了.requires_grad=True的张量的历史记录。

print(x.requires_grad)
print((x**2).requires_grad)

with torch.no_grad():
    print((x**2).requires_grad)

True
True
False

如果我们想要修改 tensor 的数值，但是又不希望被 autograd 记录(即不会影响反向传播)，那么我们可以对 tensor.data 进行操作。

x = torch.ones(1, requires_grad=True)

print(x.data)  # 还是一个tensor
print(x.data.requires_grad)  # 但是已经是独立于计算图之外

y = 2 * x
x.data *= 100  # 只改变了值，不会记录在计算图，所以不会影响梯度传播

y.backward()
print(x)  # 更改data的值也会影响tensor的值
print(x.grad)

tensor([1.])
False
tensor([100.], requires_grad=True)
tensor([2.])

3 并行计算、CUDA和cuDNN简介

在利用PyTorch做深度学习的过程中，可能会遇到数据量较大无法在单块GPU上完成，或者需要提升计算速度的场景，这时就需要用到并行计算。完成本节内容时，请你确保至少安装了一个NVIDIA GPU并安装了相关的驱动。

经过本节的学习，你将收获：

并行计算的简介
CUDA简介
并行计算的三种实现方式
使用CUDA加速训练

3.1 为什么要做并行计算

深度学习的发展离不开算力的发展，GPU的出现让我们的模型可以训练的更快，更好。所以，如何充分利用GPU的性能来提高我们模型学习的效果，这一技能是我们必须要学习的。这一节，我们主要讲的就是PyTorch的并行计算。PyTorch可以在编写完模型之后，让多个GPU来参与训练，减少训练时间。你可以在命令行使用nvidia-smi命令来查看你的GPU信息和使用情况。

3.2 为什么需要CUDA

CUDA是NVIDIA提供的一种GPU并行计算框架。对于GPU本身的编程，使用的是CUDA语言来实现的。但是，在我们使用PyTorch编写深度学习代码时，使用的CUDA又是另一个意思。在PyTorch使用 CUDA表示要开始要求我们的模型或者数据开始使用GPU了。

在编写程序中，当我们使用了 .cuda() 时，其功能是让我们的模型或者数据从CPU迁移到GPU上（默认是0号GPU）当中，通过GPU开始计算。

注：

我们使用GPU时使用的是.cuda()而不是使用.gpu()。这是因为当前GPU的编程接口采用CUDA，但是市面上的GPU并不是都支持CUDA，只有部分NVIDIA的GPU才支持，AMD的GPU编程接口采用的是OpenCL，在现阶段PyTorch并不支持。
数据在GPU和CPU之间进行传递时会比较耗时，我们应当尽量避免数据的切换。
GPU运算很快，但是在使用简单的操作时，我们应该尽量使用CPU去完成。
当我们的服务器上有多个GPU，我们应该指明我们使用的GPU是哪一块，如果我们不设置的话，tensor.cuda()方法会默认将tensor保存到第一块GPU上，等价于tensor.cuda(0)，这将有可能导致爆出out of memory的错误。我们可以通过以下两种方式继续设置。

#设置在文件最开始部分
import os
os.environ["CUDA_VISIBLE_DEVICE"] = "2" # 设置默认的显卡

 CUDA_VISBLE_DEVICE=0,1 python train.py # 使用0，1两块GPU

3.3 常见的并行的方法：

3.3.1 网络结构分布到不同的设备中(Network partitioning)

在刚开始做模型并行的时候，这个方案使用的比较多。其中主要的思路是，将一个模型的各个部分拆分，然后将不同的部分放入到GPU来做不同任务的计算。其架构如下：

这里遇到的问题就是，不同模型组件在不同的GPU上时，GPU之间的传输就很重要，对于GPU之间的通信是一个考验。但是GPU的通信在这种密集任务中很难办到，所以这个方式慢慢淡出了视野。

3.3.2 同一层的任务分布到不同数据中(Layer-wise partitioning)

第二种方式就是，同一层的模型做一个拆分，让不同的GPU去训练同一层模型的部分任务。其架构如下：

这样可以保证在不同组件之间传输的问题（这块不太理解，真的解决了么？不还是要通信么？），但是在我们需要大量的训练，同步任务加重的情况下，会出现和第一种方式一样的问题。

3.3.3 不同的数据分布到不同的设备中，执行相同的任务(Data parallelism)

第三种方式有点不一样，它的逻辑是，我不再拆分模型，我训练的时候模型都是一整个模型。但是我将输入的数据拆分。所谓的拆分数据就是，同一个模型在不同GPU中训练一部分数据，然后再分别计算一部分数据之后，只需要将输出的数据做一个汇总，然后再反传。其架构如下：

这种方式可以解决之前模式遇到的通讯问题。现在的主流方式是数据并行的方式(Data parallelism)

3.4 使用CUDA加速训练

3.4.1 单卡训练

在PyTorch框架下，CUDA的使用变得非常简单，我们只需要显式的将数据和模型通过.cuda()方法转移到GPU上就可加速我们的训练。如下：

model = Net()
model.cuda() # 模型显示转移到CUDA上

for image,label in dataloader:
    # 图像和标签显示转移到CUDA上
    image = image.cuda() 
    label = label.cuda()

3.4.2 多卡训练

PyTorch提供了两种多卡训练的方式，分别为DataParallel和DistributedDataParallel（以下我们分别简称为DP和DDP）。这两种方法中官方更推荐我们使用DDP，因为它的性能更好。但是DDP的使用比较复杂，而DP经需要改变几行代码既可以实现，所以我们这里先介绍DP，再介绍DDP。

单机多卡DP

首先我们来看单机多卡DP，通常使用一种叫做数据并行 (Data parallelism) 的策略，即将计算任务划分成多个子任务并在多个GPU卡上同时执行这些子任务。主要使用到了nn.DataParallel函数，它的使用非常简单，一般我们只需要加几行代码即可实现

model = Net()
model.cuda() # 模型显示转移到CUDA上

if torch.cuda.device_count() > 1: # 含有多张GPU的卡
	model = nn.DataParallel(model) # 单机多卡DP训练

除此之外，我们也可以指定GPU进行并行训练，一般有两种方式

nn.DataParallel函数传入device_ids参数，可以指定了使用的GPU编号

model = nn.DataParallel(model, device_ids=[0,1]) # 使用第0和第1张卡进行并行训练

要手动指定对程序可见的GPU设备

os.environ["CUDA_VISIBLE_DEVICES"] = "1,2"

多级多卡DDP

不过通过DP进行分布式多卡训练的方式容易造成负载不均衡，有可能第一块GPU显存占用更多，因为输出默认都会被gather到第一块GPU上。为此Pytorch也提供了torch.nn.parallel.DistributedDataParallel（DDP）方法来解决这个问题。

针对每个GPU，启动一个进程，然后这些进程在最开始的时候会保持一致（模型的初始化参数也一致，每个进程拥有自己的优化器），同时在更新模型的时候，梯度传播也是完全一致的，这样就可以保证任何一个GPU上面的模型参数就是完全一致的，所以这样就不会出现DataParallel那样显存不均衡的问题。不过相对应的，会比较麻烦，接下来介绍一下多机多卡DDP的使用方法。

开始之前需要先熟悉几个概念，这些还是有必要提一下的

进程组的相关概念

GROUP：进程组，默认情况下，只有一个组，一个 job 即为一个组，也即一个 world。（当需要进行更加精细的通信时，可以通过 new_group 接口，使用 world 的子集，创建新组，用于集体通信等。）
WORLD_SIZE：表示全局进程个数。如果是多机多卡就表示机器数量，如果是单机多卡就表示 GPU 数量。
RANK：表示进程序号，用于进程间通讯，表征进程优先级。rank = 0 的主机为 master 节点。如果是多机多卡就表示对应第几台机器，如果是单机多卡，由于一个进程内就只有一个 GPU，所以 rank 也就表示第几块 GPU。
LOCAL_RANK：表示进程内，GPU 编号，非显式参数，由 torch.distributed.launch 内部指定。例如，多机多卡中 rank = 3，local_rank = 0 表示第 3 个进程内的第 1 块 GPU。

DDP的基本用法 (代码编写流程)

在使用 distributed 包的任何其他函数之前，需要使用 init_process_group 初始化进程组，同时初始化 distributed 包。
使用 torch.nn.parallel.DistributedDataParallel 创建 分布式模型 DDP(model, device_ids=device_ids)
使用 torch.utils.data.distributed.DistributedSampler 创建 DataLoader
使用启动工具 torch.distributed.launch 在每个主机上执行一次脚本，开始训练

首先是对代码进行修改，添加参数 --local_rank

import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int) # 这个参数很重要
args = parser.parse_args()

这里的local_rank参数，可以理解为torch.distributed.launch在给一个GPU创建进程的时候，给这个进程提供的GPU号，这个是程序自动给的，不需要手动在命令行中指定这个参数。

local_rank = int(os.environ["LOCAL_RANK"]) #也可以自动获取

然后在所有和GPU相关代码的前面添加如下代码，如果不写这句代码，所有的进程都默认在你使用CUDA_VISIBLE_DEVICES参数设定的0号GPU上面启动

torch.cuda.set_device(args.local_rank) # 调整计算的位置

接下来我们得初始化backend，也就是俗称的后端，pytorch介绍了以下后端：

可以看到，提供了gloo，nccl，mpi，那么如何进行选择呢，官网中也给了以下建议

经验之谈
- 如果是使用cpu的分布式计算, 建议使用gloo，因为表中可以看到 gloo对cpu的支持是最好的
- 如果使用gpu进行分布式计算, 建议使用nccl。
GPU主机
- InfiniBand连接，建议使用nccl，因为它是目前唯一支持 InfiniBand 和 GPUDirect 的后端。
- Ethernet连接，建议使用nccl，因为它的分布式GPU训练性能目前是最好的，特别是对于多进程单节点或多节点分布式训练。如果在使用 nccl时遇到任何问题，可以使用gloo 作为后备选项。（不过注意，对于 GPU，gloo 目前的运行速度比 nccl 慢。）
CPU主机
- InfiniBand连接，如果启用了IP over IB，那就使用gloo，否则使用mpi
- Ethernet连接，建议使用gloo，除非有不得已的理由使用mpi。

当后端选择好了之后, 我们需要设置一下网络接口, 因为多个主机之间肯定是使用网络进行交换, 那肯定就涉及到IP之类的, 对于nccl和gloo一般会自己寻找网络接口，不过有时候如果网卡比较多的时候，就需要自己设置，可以利用以下代码

import os
# 以下二选一, 第一个是使用gloo后端需要设置的, 第二个是使用nccl需要设置的
os.environ['GLOO_SOCKET_IFNAME'] = 'eth0'
os.environ['NCCL_SOCKET_IFNAME'] = 'eth0'

可以通过以下操作知道自己的网络接口，输入ifconfig, 然后找到自己IP地址的就是, 一般就是em0, eth0, esp2s0之类的,

从以上介绍我们可以看出，当使用GPU的时候, nccl的效率是高于gloo的，我们一般还是会选择nccl后端，设置GPU之间通信使用的后端和端口：

# ps 检查nccl是否可用
# torch.distributed.is_nccl_available ()
torch.distributed.init_process_group(backend='nccl') # 选择nccl后端，初始化进程组

之后，使用 DistributedSampler 对数据集进行划分。它能帮助我们将每个 batch 划分成几个 partition，在当前进程中只需要获取和 rank 对应的那个 partition 进行训练：

# 创建Dataloader
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=16, sampler=train_sampler)

注意： testset不用sampler

然后使用torch.nn.parallel.DistributedDataParallel包装模型：

# DDP进行训练
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank])

如何启动DDP

那么如何启动DDP，这不同于DP的方式，需要使用torch.distributed.launch启动器，对于单机多卡的情况：

CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch --nproc_per_node=4 main.py
# nproc_per_node: 这个参数是指你使用这台服务器上面的几张显卡

有时候虽然说，可以简单使用DP，但是DDP的效率是比DP高的，所以很多时候单机多卡的情况，我们还是会去使用DDP

3.4.3 DP与DDP的优缺点

DP 的优势

nn.DataParallel没有改变模型的输入输出，因此其他部分的代码不需要做任何更改，非常方便，一行代码即可搞定。

DP 的缺点

DP进行分布式多卡训练的方式容易造成负载不均衡，第一块GPU显存占用更多，因为输出默认都会被gather到第一块GPU上，也就是后续的loss计算只会在cuda:0上进行，没法并行。

除此之外DP只能在单机上使用，且DP是单进程多线程的实现方式，比DDP多进程多线程的方式会效率低一些。

DDP的优势

1. 每个进程对应一个独立的训练过程，且只对梯度等少量数据进行信息交换。

DDP 在每次迭代中，每个进程具有自己的 optimizer ，并独立完成所有的优化步骤，进程内与一般的训练无异。

在各进程梯度计算完成之后，各进程需要将梯度进行汇总平均，然后再由 rank=0 的进程，将其 broadcast 到所有进程。之后，各进程用该梯度来独立的更新参数。而 DP是梯度汇总到主 GPU，反向传播更新参数，再广播参数给其他的 GPU。

DDP 中由于各进程中的模型，初始参数一致 (初始时刻进行一次 broadcast)，而每次用于更新参数的梯度也一致，因此，各进程的模型参数始终保持一致。

而在DP 中，全程维护一个 optimizer，对各 GPU 上梯度进行求和，而在主 GPU 进行参数更新，之后再将模型参数 broadcast 到其他 GPU。

相较于DP，DDP传输的数据量更少，因此速度更快，效率更高。

2. 每个进程包含独立的解释器和 GIL。

一般使用的 Python 解释器 CPython：是用 C 语言实现 Pyhon，是目前应用最广泛的解释器。全局锁使 Python 在多线程效能上表现不佳，全局解释器锁（Global Interpreter Lock）是 Python 用于同步线程的工具，使得任何时刻仅有一个线程在执行。

由于每个进程拥有独立的解释器和 GIL，消除了来自单个 Python 进程中的多个执行线程，模型副本或 GPU 的额外解释器开销和 GIL-thrashing ，因此可以减少解释器和 GIL 使用冲突。这对于严重依赖 Python runtime 的 models 而言，比如说包含 RNN 层或大量小组件的 models 而言，这尤为重要。

DDP 的缺点

暂时来说，DDP是采用多进程多线程的方式，并且训练速度较高，他的缺点主要就是，需要修改比较多的代码，比DP的一行代码较为繁琐许多。

3.5 通俗解释单机单卡、单机多卡和多机多卡

这里用一个打工搬砖的比喻来通俗解释 单机单卡、单机多卡 和 多机多卡：

3.5.1 单机单卡：一个工人、一辆车

比喻：只有一个工人（单块 GPU），用一辆手推车（计算资源）搬砖（计算任务）。
特点：
- 一个人搬砖效率有限，但全程由一个人负责，任务分配和管理简单。
- 适合小任务，比如小房子装修。
实际应用：
- 常用于小规模深度学习任务或单机单模型运行。

3.5.2 单机多卡：一个工地、多个工人，每人一辆车

比喻：一个工地里有多个人（多块 GPU），每个人都有自己的手推车，大家一起搬砖。
特点：
- 每个工人可以负责一部分砖（数据并行）或者不同的任务（模型并行）。
- 工人们在一个工地上，距离近（同一台机器），交流方便（高速显存/总线通信）。
- 如果需要每个人把搬的砖同步起来，就会有一定的通信时间（同步开销）。
- 适合中等规模任务，比如盖个小楼。
实际应用：
- 深度学习中使用多块 GPU 同时训练模型（如训练 ResNet）。

3.5.3 多机多卡：多个工地、多个工人，每人一辆车

比喻：现在任务太大，一个工地不够用了，于是开了多个工地，每个工地上都有工人（每台机器配多个 GPU），每个工人还是推着自己的车。
特点：
- 工地之间需要通信（通过网络传输数据），效率会受到网络速度影响（比如用 PCIe、NVLink、InfiniBand）。
- 工人间可以分开搬砖，但要协作完成整个任务，通信和协调的复杂性增加。
- 适合超大规模任务，比如建高楼大厦。
实际应用：
- 分布式训练，像 GPT-4 这种需要同时用几十到上百块 GPU 的任务。

3.5.4 对比总结

3.6 参考资料

4 AI硬件加速设备

4.1 CPU和GPU

在进行模型部署和训练时，我们有时会受限于CPU和GPU的性能。这时，专用的AI芯片就显得尤为重要。在正式开始本节内容之前，我们先了解一下什么是CPU和GPU。

CPU即Central Processing Unit，中文名为中央处理器，是我们电脑中的核心配件。它的功能主要是处理指令、执行操作、控制时间、处理数据。

在现代计算机体系结构中，CPU 对计算机的所有硬件资源（如存储器、输入输出单元）进行控制调配、执行通用运算的核心硬件单元。CPU 是计算机的运算和控制核心。计算机系统中所有软件层的操作，最终都将通过指令集映射为CPU的操作。

GPU即Graphics Processing Unit，中文名为图形处理单元。在传统的冯·诺依曼结构中，CPU 每执行一条指令都需要从存储器中读取数据，根据指令对数据进行相应的操作。从这个特点可以看出，CPU 的主要职责并不只是数据运算，还需要执行存储读取、指令分析、分支跳转等命令。深度学习算法通常需要进行海量的数据处理，用 CPU 执行算法时，CPU 将花费大量的时间在数据/指令的读取分析上，而 CPU的频率、内存的带宽等条件又不可能无限制提高，因此限制了处理器的性能。而 GPU 的控制相对简单，大部分的晶体管可以组成各类专用电路、多条流水线，使得 GPU 的计算速度远高于CPU；同时 GPU 拥有了更加强大的浮点运算能力，可以缓解深度学习算法的训练难题，释放人工智能的潜能。

需要注意的是，GPU没有独立工作的能力，必须由CPU进行控制调用才能工作，且GPU的功耗一般比较高。因此，随着人工智能的不断发展，高功耗低效率的GPU不再能满足AI训练的要求，为此，一大批功能相对单一，但速度更快的专用集成电路相继问世。接下来我们了解一下什么是专用集成电路：

专用集成电路（Application-Specific Integrated Circuit，ASIC）是专用定制芯片，即为实现特定要求而定制的芯片。定制的特性有助于提高 ASIC 的性能功耗比。ASIC的缺点是电路设计需要定制，相对开发周期长，功能难以扩展。但在功耗、可靠性、集成度等方面都有优势，尤其在要求高性能、低功耗的移动应用端体现明显。下文提到的谷歌的TPU，寒武纪的NPU都属于ASIC的范畴。

下面让我们进入本节的内容，经过本节的学习，你将收获：

什么是TPU
什么是NPU

4.2 TPU

TPU即Tensor Processing Unit，中文名为张量处理器。2006年，谷歌开始计划为神经网络构建一个专用的集成电路（ASIC）。随着计算需求和数据量的不断上涨，这个需求在2013年开始变得尤为紧迫。于是，谷歌在2015年6月的IO开发者大会上推出了为优化自身的TensorFlow框架而设计打造的一款计算神经网络专用芯片。它主要用于进行搜索，图像，语音等模型和技术的处理。

截至目前，谷歌已经发行了四代TPU芯片。

4.2.1 芯片架构设计

TPU的设计架构如下图

上图：In-datacenter performance analysis of a tensor processing unit，figure 1

由上图可见，整个TPU中最重要的计算单元是右上角黄色的矩阵乘单元“Matrix Multiply Unit”，它包含256x256个MAC部件，每一个能够执行有符号或者无符号的8位乘加操作。它的输入为权重数据队列FIFO和统一缓冲Unified Buffer，即图中指向它的两个蓝色部分。在计算结束后，16位结果被收集并传递到位于矩阵单元下方的4MiB 32位蓝色累加器Accumulators中，之后由黄色的激活单元在累加后执行非线性函数，并最终将数据返回给统一缓冲。

Matrix Multiply Unit矩阵处理器作为TPU的核心部分，它可以在单个时钟周期内处理数十万次矩阵（Matrix）运算。MMU有着与传统CPU、GPU截然不同的架构，称为脉动阵列（systolic array）。之所以叫“脉动”，是因为在这种结构中，数据一波一波地流过芯片，与心脏跳动供血的方式类似。而如下图所示，CPU和GPU在每次运算中都需要从多个寄存器（register）中进行存取，而TPU的脉动阵列将多个运算逻辑单元（ALU）串联在一起，复用从一个寄存器中读取的结果。每个ALU单元结构简单，一般只包含乘法器、加法器以及寄存器三部分，适合大量堆砌。

但是，在极大增加数据复用、降低内存带宽压力的同时，脉动阵列也有两个缺点，即数据重排和规模适配。第一，脉动矩阵主要实现向量/矩阵乘法。以CNN计算为例，CNN数据进入脉动阵列需要调整好形式，并且严格遵循时钟节拍和空间顺序输入。数据重排的额外操作增加了复杂性。第二，在数据流经整个阵列后，才能输出结果。当计算的向量中元素过少，脉动阵列规模过大时，不仅难以将阵列中的每个单元都利用起来，数据的导入和导出延时也随着尺寸扩大而增加，降低了计算效率。因此在确定脉动阵列的规模时，在考虑面积、能耗、峰值计算能力的同时，还要考虑典型应用下的效率。

4.2.2 技术特点

AI加速专用

TPU的架构属于Domain-specific Architecture，也就是特定领域架构。它的定位准确，架构简单，单线程控制，定制指令集使得它在深度学习运算方面效率极高，且容易扩展。相比之下，传统诸如CPU、GPU等通用处理器必须考虑灵活性和兼容性，有太重的包袱。但TPU这种特点也决定它只能被限制用于深度学习加速场景。

脉动阵列设计

TPU采用了与传统CPU和GPU截然不同的脉动阵列（systolic array）结构来加速AI运算，脉动阵列能够在一个时钟周期内处理数十万次矩阵运算，在每次运算过程中，TPU能够将多个运算逻辑单元（ALU）串联在一起，并复用从一个寄存器中取得的结果。这种设计，不仅能够将数据复用实现最大化，减少芯片在运算过程中的内存访问次数，提高AI计算效率，同时也降低了内存带宽压力，进而降低内存访问的能耗。

MMU的脉动阵列包含256 × 256 = 65,536个ALU，也就是说TPU每个周期可以处理65,536次8位整数的乘法和加法。

TPU以700兆赫兹的功率运行，也就是说，它每秒可以运行65,536 × 700,000,000 = 46 × 1012次乘法和加法运算，或每秒92万亿（92 × 1012）次矩阵单元中的运算。

上图：In-datacenter performance analysis of a tensor processing unit，figure 4

确定性功能和大规模片上内存

如图是TPU的平面设计简图，黄色为MMU运算单元，蓝色是统一缓存和累加器等数据单元，绿色是I/O，红色是逻辑控制单元。

上图：In-datacenter performance analysis of a tensor processing unit，figure 2

传统GPU由于片上内存较少，因此在运行过程中需要不断地去访问片外动态随机存取存储器（DRAM），从而在一定程度上浪费了不必要的能耗。与CPU和GPU相比，TPU的控制单元更小，更容易设计，面积只占了整个冲模的2%，给片上存储器和运算单元留下了更大的空间。如上图所示的TPU一代架构中，总共设计了占总芯片面积35%的内存，其中包括24MB的局部内存、4MB的累加器内存，以及用于与主控处理器对接的内存。这一比例大大超出了GPU等通用处理器，节约了大量片外数据访存能耗，使得TPU计算的能效比大大提高。从TPU二代开始采用HBM片上高带宽内存，虽然和最新一代GPU片上内存技术相同，但是TPU芯片的面积要远远小于GPU。硅片越小，成本越低，良品率也越高。

另外，由于TPU是一个单用途芯片，不需要考虑缓存、分支预测、多道处理等问题。这就意味着TPU的功能是单一且确定的。因此，我们可以使用TPU轻易的预测运行一个神经网络需要多长时间，这样我们就能让芯片以吞吐量接近峰值的状态运行，同时严格控制延迟。

4.3 NPU

NPU即Neural-network Processing Unit，中文名为神经网络处理器，它采用“数据驱动并行计算”的架构，特别擅长处理视频、图像类的海量多媒体数据。

长期以来，应用需求一直牵动着嵌入式技术的发展方向。随着深度学习神经网络的兴起，人工智能、大数据时代的来临，CPU和GPU渐渐难以满足深度学习的需要，面对日渐旺盛的需求和广大的预期市场，设计一款专门用于神经网络深度学习的高效智能处理器显得十分必要，因此NPU应运而生。

从技术角度看，深度学习实际上是一类多层大规模人工神经网络。它模仿生物神经网络而构建，由若干人工神经元结点互联而成。神经元之间通过突触两两连接，突触记录了神经元间联系的权值强弱。由于深度学习的基本操作是神经元和突触的处理，神经网络中存储和处理是一体化的，都是通过突触权重来体现，而冯·诺伊曼结构中，存储和处理是分离的，分别由存储器和运算器来实现，二者之间存在巨大的差异。当用现有的基于冯·诺伊曼结构的经典计算机(如X86处理器和英伟达GPU)运行神经网络应用时，就不可避免地受到存储和处理分离式结构的制约，因而影响效率。因此，专门针对人工智能的专业芯片NPU更有研发的必要和需求。

在NPU的设计上，中国走在了世界前列。下面我们将以寒武纪的DianNao系列架构为例，来简要介绍NPU。

4.3.1 DianNao

上图：DianNao: a small-footprint high-throughput accelerator for ubiquitous machine-learning，figure 9

基于神经网络的人工智能算法，是模拟人类大脑内部神经元的结构。上图中的neuron代表的就是单个神经元，synapse代表神经元的突触。这个模型的工作模式，就要结合高中生物课的知识了。

一个神经元，有许多突触，给别的神经元传递信息。同样，这个神经元，也会接收来自许多其他神经元的信息。这个神经元所有接受到的信息累加，会有一个强烈程度，在生物上是以化学成分的形式存在，当这些信息达到一定的强烈程度，就会使整个神经元处于兴奋状态（激活），否则就是不兴奋（不激活）。如果兴奋了，就给其他神经元传递信息，如果不兴奋，就不传递。这就是单独一个神经元的工作模式。那么有成千上万个这样的神经元组合起来，就是一个神经网络模型。

那么DianNao是如何模拟神经元进行工作的呢，我们可以看看它的内部结构图：

上图：DianNao: a small-footprint high-throughput accelerator for ubiquitous machine-learning，figure 11

如图所示，上图中浅蓝色的部分就是用硬件逻辑模拟的神经网络架构，称为NFU（Neural Functional Units）。它可以被细分为三个部分，即途中的NFU-1，NFU-2，和NFU-3。

NFU-1是乘法单元，它采用16bit定点数乘法器，1位符号位，5位整数位，10位小数位。该部分总计有256个乘法器。这些乘法器的计算是同时的，也就是说，在一个周期内可以执行256次乘法。

NFU-2是加法树，总计16个，每一个加法树都是8-4-2-1这样的组成结构，即就是每一个加法树中都有15个加法器。

NFU-3是非线性激活函数，该部分由分段线性近似实现非线性函数，根据前面两个单元计算得到的刺激量，从而判断是否需要激活操作。

当需要实现向量相乘和卷积运算时，使用NFU-1完成对应位置元素相乘，NFU-2完成相乘结果相加，最后由NFU-3完成激活函数映射。完成池化运算时，使用NFU-2完成多个元素取最大值或取平均值运算。由此分析，尽管该运算模块非常简单，也覆盖了神经网络所需要的大部分运算。

4.3.2 DaDianNao

作为DianNao的多核升级版本，DaDianNao的运算单元NFU与DianNao基本相同，最大的区别是为了完成训练任务多加了几条数据通路，且配置更加灵活。NFU的尺寸为16x16，即16个输出神经元，每个输出神经元有16个输入（输入端需要一次提供256个数据）。同时，NFU可以可选的跳过一些步骤以达到灵活可配置的功能。DaDianNao的NFU结构如下所示：

上图：DaDianNao: A Machine-Learning Supercomputer，figure 6

4.3.3 ShiDianNao

ShiDianNao是机器视觉专用加速器，集成了视频处理的部分，它也是DianNao系列中唯一一个考虑运算单元级数据重用的加速器，也是唯一使用二维运算阵列的加速器，其加速器的运算阵列结构如下所示：

上图：ShiDianNao: Shifting vision processing closer to the sensor，figure 5

ShiDianNao的运算阵列为2D格点结构，对于每一个运算单元（节点）而言，运算所使用的参数统一来源于Kernel，而参与运算的数据则可能来自于：数据缓存NBin，下方的节点，右侧的节点。

下图为每个运算单元的结构：

上图：ShiDianNao: Shifting vision processing closer to the sensor，figure 6

该计算节点的功能包括转发数据和进行计算：

转发数据：每个数据可来源于右侧节点，下方节点和NBin，根据控制信号选择其中一个存储到输入寄存器中，且根据控制信号可选的将其存储到FIFO-H和FIFO-V中。同时根据控制信号选择FIFO-H和FIFO-V中的信号从FIFO output端口输出

进行计算：根据控制信号进行计算，包括相加，累加，乘加和比较等，并将结果存储到输出寄存器中，并根据控制信号选择寄存器或计算结果输出到PE output端口。

对于计算功能，根据上文的结构图，可以发现，PE支持的运算有：kernel和输入数据相乘并与输出寄存器数据相加（乘加），输入数据与输出寄存器数据取最大或最小（应用于池化），kernel与输入数据相加（向量加法），输入数据与输出寄存器数据相加（累加）等。

4.3.4 PuDianNao

作为DianNao系列的收山之作，PuDianNao的运算单元是电脑系列中唯一一个异构的，除了有MLU（机器学习单元）外，还有一个ALU用于处理通用运算和MLU无法处理的运算，其运算单元（上）和MLU（下）结构如下图所示：

上图：PuDianNao: A Polyvalent Machine Learning Accelerator，figure 11&12

该MLU共分为6层：

计数层/比较层：这一层的处理为两个数按位与或比较大小，结果将被累加，这一层可以单独输出且可以被bypass

加法层：这一层为两个输入对应相加，这一层可以单独输出且可以被bypass

乘法层：这一层为两个输入或上一层（加法层）结果对应位置相乘，可以单独输出

加法树层：将乘法层的结果累加

累加层：将上一层（加法树层）的结果累加，可以单独输出

特殊处理层：由一个分段线性逼近实现的非线性函数和k排序器（输出上一层输出中最小的输出）组成

该运算单元是DianNao系列中功能最多的单元，配置非常灵活。例如实现向量相乘（对应位置相乘后累加）时，弃用计数层，加法层，将数据从乘法层，加法树层和累加层流过即可实现。

PuDianNao支持7种机器学习算法：神经网络，线性模型，支持向量机，决策树，朴素贝叶斯，K临近和K类聚，所需要支持的运算较多，因此PuDianNao的运算分析主要集中在存储方面，其运算核心的设计中说明PuDianNao支持的运算主要有：向量点乘，距离计算，计数，排序和非线性函数。其他未覆盖的计算使用ALU实现。