Pytorch从入门到精通

时间：12-08来源：作者：点击数：31

记得刚开始学TensorFlow的时候，那给我折磨的呀，我一直在想这个TensorFlow官方为什么搭建个网络还要画什么静态图呢，把简单的事情弄得麻烦死了，直到这几天我开始接触Pytorch，发现Pytorch是就是不用搭建静态图的Tensorflow版本，就想在用numpy一样，并且封装了很多深度学习高级API，numpy数据和Tensor数据相互转换不用搭建会话了，只需要一个转换函数，搭建起了numpy和TensorFlow爱的桥梁。

Pytorch自17年推出以来，一度有赶超TensorFlow的趋势，是因为Pytorch采用动态图机制，替代Numpy使用GPU的功能，搭建网络灵活。

Pytorch和TensorFlow的区别：

TensorFlow是基于静态计算图的，静态计算图是先定义后运行，一次定义多次运行（Tensorflow 2.0也开始使用动态计算图）
PyTorch是基于动态图的，是在运行的过程中被定义的，在运行的时候构建，可以多次构建多次运行

上手难度：tensorflow 1 > tensorflow 2 > pytorch

工业界：tensorflow 1>tensorflow 2 > pytorch

学术界：pytorch > tesnroflow 2 > tensorflow 1(已经被谷歌抛弃)

pytorch的数据类型

Data type	CPU tensor	GPU tensor
torch.float32	torch.FloatTensor	torch.cuda.FloatTensor
torch.floar64	torch.DoubleTensor	torch.cuda.DoubleTensor
torch.int32	torch.IntTensor	torch.cuda.IntTensor
torch.int64	torch.LongTensor	torch.cuda.LongTensor

即便是同一个变量同时部署在CPU和GPU上面是不一样的

这篇文章的所有代码，请务必手敲！！！

安装

这个网址包含pytorch与cuda的对应关系

由于我的cuda是 10.0

我选择的安装命令是：

CPU版本：pip install torch==1.4.0 torchvision==0.5.0 -f https://download.pytorch.org/whl/torch_stable.html

GPU版本：pip install torch==1.4.0+cu92 torchvision==0.5.0+cu92 -f https://download.pytorch.org/whl/torch_stable.html

官网给的安装命令是：

pip install torch==1.2.0 torchvision==0.4.0
# 或
pip install torch==1.2.0+cu92 torchvision==0.4.0+cu92 -f https://download.pytorch.org/whl/torch_stable.html

我原本是：torch-1.2.0+cu92

目前我装的是（安装torchaudio时候帮我升级的）：torch-1.6.0+cu101 torchaudio-0.6.0

创建张量

torch的数据类型torch.float32、torch.floar64、torch.float16、torch.int8、torch.int16、torch.int32、torch.int64。当数据在GPU上时，数据类型需要加上cuda，例：torch.cuda.FloatTensor

tensor.shape/tensor.size()：获取张量的shape

tensor.reshape()/tensor.view()：修改张量的shape

tensor.item()：如果我们的张量只有一个数值，可以使用.item()获取，常用于获取loss值

tensor.dim()：返回张量的阶数\维度

tensor.type()：查看数据类型

tensor.cuda()：将tensor转换到GPU上

# 0阶\0维 张量
a = torch.tensor(1.3)
print(a)  # tensor(1.3000)
print(a.shape)  # torch.Size([])

# 1阶张量
a = torch.tensor([1.1])
b = torch.tensor([1.1, 2.2])
print(a.shape)  # torch.Size([1])
print(b.shape)  # torch.Size([2])

# 2阶张量
a = torch.randn(2,3)
print(a.shape)  # torch.Size([2, 3])

# 3阶张量
a = torch.rand(1, 2, 3)
print(a.shape)  # torch.Size([1, 2, 3])

直接指定tensor的数值

print(torch.tensor([2.,3.2]))
# tensor([2.0000, 3.2000])
print(torch.FloatTensor([2.,3.2]))
# tensor([2.0000, 3.2000])
print(torch.tensor([[2.,3.2],[1.,22.3]]))
# tensor([[ 2.0000,  3.2000],
#         [ 1.0000, 22.3000]])

定义未初始化张量

print(torch.empty(2,3))
# tensor([[2.5657e-05, 6.3199e-43, 2.5657e-05],
#         [6.3199e-43, 2.5855e-05, 6.3199e-43]])
print(torch.FloatTensor(2,3))
# tensor([[2.5804e-05, 6.3199e-43, 8.4078e-45],
#         [0.0000e+00, 1.4013e-45, 0.0000e+00]])
print(torch.IntTensor(2,3))
# tensor([[937482688,       451,         1],
#         [        0,         1,         0]], dtype=torch.int32)

设置tensor数据的默认类型type

print(torch.tensor([1.2,3]).type())

torch.set_default_tensor_type(torch.DoubleTensor)
print(torch.tensor([1.3,3]).type())

torch.ones(size)/zero(size)/eye(size)：返回全为1/0/单位对角张量

torch.full(size, fill_value)：返回以size大小填充fill_value的张量

torch.rand(size)：返回[0, 1)之间的均匀分布张量

torch.randn(size)：均值为0，方差为1的正态分布

torch.*_like(input)：返回一个和input shape一样的张量，*可以为rand、randn...

torch.randint(low=0, high, size)：返回shape=size，[low, high)之间的随机整数

torch.arange()：和np.arange类似用法

torch.linspace(start, end, step=1000)：返回start和end之间等距steps点的一维步长张量。

torch.logspace(start, end, steps=1000, base=10.0)：返回 $base^{start}$ 和 $base^{end}$ 之间等距steps点的一维步长张量。

torch.randperm(n)：返回从0到n-1的整数的随机排列

b = torch.rand(4)
idx = torch.randperm(4)
print(b)    # tensor([0.0224, 0.7826, 0.5529, 0.2261])
print(idx)  # tensor([0, 2, 1, 3])
print(b[idx])   # tensor([0.5573, 0.6121, 0.6581, 0.1892])

pytorch与numpy变量转换

numpy变量 -----> torch变量：torch.from_numpy(ndarray)

torch变量 ------> numpy变量：tensor.numpy()

x = np.array([[1, 2], [3, 4]])
y = torch.from_numpy(x)  # 转换为 torch数据
z = y.numpy()  # 转换为 numpy 数据

维度变换

tensor.reshape()：维度变换
tensor.view()：维度变换
squeeze：去除对应维度为1的维度
unsqueeze：往对应位置索引插入一个维度

a = torch.rand(4, 1, 28, 28)
b = a.unsqueeze(0)
print(b.shape)  # torch.Size([1, 4, 1, 28, 28])

a = torch.rand(4, 1, 1, 28)
b = a.squeeze()
print(b.shape)  # torch.Size([4, 28])
b = a.squeeze(1)
print(b.shape)  # torch.Size([4, 1, 28])

tensor.expand(*size)：返回具有单个尺寸扩展到更大尺寸的张量
tensor.repeat(*size)：沿指定尺寸重复此张量

x = torch.tensor([[1], [2], [3]])
print(x.size())  # torch.Size([3, 1])
print(x.expand(3, 4))
# tensor([[ 1,  1,  1,  1],
#         [ 2,  2,  2,  2],
#         [ 3,  3,  3,  3]])
print(x.expand(-1, 4))  # -1表示不改变维度的大小
# tensor([[ 1,  1,  1,  1],
#         [ 2,  2,  2,  2],
#         [ 3,  3,  3,  3]])

b = torch.tensor([1, 2, 3])  # torch.Size([1, 3])
print(b.repeat(4, 2).shape)  # torch.Size([4, 6])
print(b.repeat(4, 2, 1).size())  # torch.Size([4, 2, 3])

tensor.transpose()：调换张量指定维度的顺序
tensor.permute()：将张量按指定顺序排列

b = torch.rand(4, 3, 28, 32)

print(b.transpose(1, 3).shape)  # torch.Size([4, 32, 28, 3])
print(b.permute(0, 2, 3, 1).shape)  # torch.Size([4, 28, 32, 3])

torch.cat(inputs, dim=0) ：在给定维度上对输入的张量进行连接拼接
torch.stack()：沿着一个新维度对输入张量序列进行拼接

import torch

a = torch.randn(2, 3)
b = torch.randn(2, 3)

c = torch.cat((a, b), 0).size()    # (4, 3)
d = torch.stack((a, b), 0).size()    # (2, 2, 3)

索引与切片操作

a = torch.rand(4, 3, 28, 28)
print(a.shape)  # torch.Size([4, 3, 28, 28])

# 索引
print(a[0, 0].shape)  # torch.Size([28, 28])
print(a[0, 0, 2, 4])  # tensor(0.1152)

# 切片
print(a[:2].shape)  # torch.Size([2, 3, 28, 28])
print(a[:2, :2, :, :].shape)  # torch.Size([2, 2, 28, 28])
print(a[:2, -1:, :, :].shape)  # torch.Size([2, 1, 28, 28])

# ...的用法
print(a[...].shape)  # torch.Size([4, 3, 28, 28])
print(a[0, ...].shape)  # torch.Size([3, 28, 28])
print(a[:, 1, ...].shape)  # torch.Size([4, 28, 28])
print(a[..., :2].shape)  # torch.Size([4, 3, 28, 2])

掩码取值

x = torch.tensor([[0.4573, 0.8571], [0.7237, 0.0265]])

mask = x.gt(0.5)  # ge 是符号 >。满足条件返回True,否则False
print(mask)
# tensor([[False,  True],
#         [ True, False]])

mask = x > 0.5  # 满足条件返回True,否则False
# tensor([[False,  True],
#         [ True, False]])

x_index = x[x > 0.5]  # 把大于0.5的值都取出来
# tensor([0.8571, 0.7237])

print(torch.masked_select(x, mask))  # 返回满足条件的值
# tensor([0.8571, 0.7237])

mask = torch.where(x > 0.5)  # 返回满足条件的索引((行,行), (列,列))
# (tensor([0, 1]), tensor([1, 0]))

mask = torch.nonzero(x > 0.5)
# tensor([[0, 1],   # 第0行第1列
#         [1, 0]])  # 第1行第0列

通过torch.take取值

src = torch.tensor([[4,3,5],[6,7,8]])
print(torch.take(src, torch.tensor([0,2,5])))# 取第0、2、5个元素
# tensor([4, 5, 8])

数学运算

加法：tensor1 + tensor2 或 torch.add(tensor1, tensor2)

减法：tensor1 - tensor2 或 torch.sub(tensor1, tensor2)

a = torch.rand(3, 4)  # dim=2
b = torch.rand(4)     # dim=1
# 加法
print(a + b)
print(torch.add(a, b))
# 减法
print(a - b)
print(torch.sub(a, b))

乘法

tensor1 * tensor2 ：对位元素相乘

torch.mul(tensor1, other) ：input是矩阵，other可以是矩阵或标量，是矩阵时，对位相乘，就可以广播

torch.mm(tensor1, tensor2) ：只能处理二维矩阵的乘法

tensor1 @ tensor2 ：二维矩阵相乘

torch.bmm(tensor1, tensor2) ：在torch.mm的基础上加了个batch计算，不能广播

torch.matmul(input, other) ：适用性最多的，能处理batch，能广播的矩阵

如果第一个参数是一维，第二个是二维，那么给第一个提供一个维度
如果第一个是二维，第二个是一维，就是矩阵乘向量
带有batch的情况，可保留batch计算
维度不同时，可先广播，再batch计算

“广播”注释：Broadcasting，在运算中，不同大小的两个 array 应该怎样处理的操作。通常情况下，小一点的数组会被 broadcast 到大一点的，这样才能保持大小一致。Broadcasting 过程中的循环操作都在 C 底层进行，所以速度比较快。但也有一些情况下 Broadcasting 会带来性能上的下降。

总结：对位相乘用 torch.mul ，二维矩阵乘法用 torch.mm ，batch二维矩阵用 torch.bmm ，batch、广播用 torch.matmul

平方

a = torch.full([2,2], 2)    # 创建一个shape=[2,2]值为2的数组
print(a.pow(2))
print(a**2)

平方根

a = torch.full([2,2], 4)    # 创建一个shape=[2,2]值为4的数组
print(a.sqrt())    # 平方根
print(a**(0.5))

torch.exp()：e的指数冥

torch.log()：取对数

tensor.floor()：向下取整

tensor.ceil() ：向上取整

tensor.round()：四舍五入

tensor.trunc()：取整数值

tensor.frac()：取小数值

tensor.clamp(min,max)：不足最小值的变成最小值，大于最大值的变成最大值

torch.mean()：求均值

torch.sum()：求和

torch.max\torch.min：求最大最小值

torch.prod(input, dtype=None) ：返回input中所有元素的乘积

torch.argmin(input)\torch.argmax(input)：返回input张量中所有元素的最小值\最大值的索引

torch.where(condition, x, y)：如果符合条件返回x，如果不符合条件返回y

torch.gather(input, dim, index)：沿dim指定的轴收集值

input：输入tensor
dim：索引所沿的轴
index：要收集的元素的索引

t = torch.tensor([[1,2],[3,4]])
torch.gather(t, 1, torch.tensor([[0,0],[1,0]]))
# tensor([[ 1,  1],
#         [ 4,  3]])

autograd：自动求导

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

让我们用一些简单的例子来看看吧。

张量

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

如果要阻止一个张量被跟踪历史，可以调用 .detach() 方法将其与计算历史分离，并阻止它未来的计算记录被跟踪。

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

每个张量都有一个 .grad_fn 属性，该属性引用了创建 Tensor 自身的Function(除非这个张量是用户手动创建的，即这个张量的 grad_fn 是 None )。

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

有一部分有点难，需要多看几遍。

https://pytorch.apachecn.org/docs/1.2/beginner/blitz/autograd_tutorial.html

如果设置torch.tensor_1(requires_grad=True)，那么会追踪所有对该张量tensor_1的所有操作。

import torch

# 创建一个张量并设置 requires_grad=True 用来追踪他的计算历史
x = torch.ones(2, 2, requires_grad=True)
print(x)
# tensor([[1., 1.],
#         [1., 1.]], requires_grad=True)

当Tensor完成一个计算过程，每个张量都会自动生成一个.grad_fn属性

# 对张量进行计算操作，grad_fn已经被自动生成了。
y = x + 2
print(y)
# tensor([[3., 3.],
#         [3., 3.]], grad_fn=<AddBackward>)
print(y.grad_fn)
# <AddBackward object at 0x00000232535FD860>

# 对y进行一个乘法操作
z = y * y * 3
out = z.mean()

print(z)
# tensor([[27., 27.],
#         [27., 27.]], grad_fn=<MulBackward>) 
print(out)
# tensor(27., grad_fn=<MeanBackward1>)

.requires_grad_(...) 可以改变张量的requires_grad属性。

import torch

a = torch.randn(2, 2)
a = ((a * 3) / (a - 1))
print(a.requires_grad)      # 默认是requires_grad = False
a.requires_grad_(True)
print(a.requires_grad)      # True
b = (a * a).sum()
print(b.grad_fn)        # <SumBackward0 object at 0x000002325360B438>

梯度

回顾到上面

import torch

# 创建一个张量并设置 requires_grad=True 用来追踪他的计算历史
x = torch.ones(2, 2, requires_grad=True)
print(x)
# tensor([[1., 1.],
#         [1., 1.]], requires_grad=True)


# 对张量进行计算操作，grad_fn已经被自动生成了。
y = x + 2
print(y)
# tensor([[3., 3.],
#         [3., 3.]], grad_fn=<AddBackward>)
print(y.grad_fn)
# <AddBackward object at 0x00000232535FD860>

# 对y进行一个乘法操作
z = y * y * 3
out = z.mean()

print(z)
# tensor([[27., 27.],
#         [27., 27.]], grad_fn=<MulBackward>)
print(out)
# tensor(27., grad_fn=<MeanBackward1>)

让我们来反向传播，运行 out.backward() ，等于out.backward(torch.tensor(1.))

对out进行反向传播， $out = \frac{1}{4}\sum_i z_i$ ，其中 $z_i = 3(x_i+2)^2$ ，因为方向传播中torch.tensor=1（out.backward中的参数）因此 $z_i\bigr\rvert_{x_i=1} = 27$

对于梯度 $\frac{\partial out}{\partial x_i} = \frac{3}{2}(x_i+2)$ ，把 $x_i=1$ 代入 $\frac{\partial out}{\partial x_i}\bigr\rvert_{x_i=1} = \frac{9}{2} = 4.5$

print(out)  # tensor(27., grad_fn=<MeanBackward1>)

print("*"*50)
out.backward()
# 打印梯度
print(x.grad)
# tensor([[4.5000, 4.5000],
#         [4.5000, 4.5000]])

对吃栗子找到规律，才能看懂

import torch

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

y = x * 2
while y.data.norm() < 1000:
    y = y * 2

print(y)  # tensor([-920.6895, -115.7301, -867.6995], grad_fn=<MulBackward>)
gradients = torch.tensor([0.1, 1.0, 0.0001], dtype=torch.float)

# 把gradients代入y的反向传播中
y.backward(gradients)   

# 计算梯度
print(x.grad)   # tensor([ 51.2000, 512.0000,   0.0512])

为了防止跟踪历史记录，可以将代码块包装在with torch.no_grad()：中。在评估模型时特别有用，因为模型的可训练参数的属性可能具有requires_grad = True，但是我们不需要梯度计算。

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

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

自定义数据集

这一节从另一篇文章截取了部分，想要具体了解请参考：pytorch加载语音类自定义数据集

pytorch内部集成了一些常用的数据集调用接口，但是当我们需要使用自己的数据集训练神经网络时，就需要自定义数据集，在pytorch中，提供了一些类，方便我们定义自己的数据集合

torch.utils.data.Dataset：所有继承他的子类都应该重写 __len()__ ， __getitem()__ 这两个方法
- __len()__ ：返回数据集中数据的数量
- __getitem()__ ：返回支持下标索引方式获取的一个数据
torch.utils.data.DataLoader：对数据集进行包装，可以设置batch_size、是否shuffle....

要创建自己的自定义的数据集首先要创建一个自定义的数据集类，这个类要继承torch.utils.data.Dataset类，并且我们还需要重写父类的__init__()和__getitem__()方法，具体形式如下

from torch.utils.data import Dataset

# 创建 MyselfDataset 数据集类
class MyselfDataset(Dataset):
    def __init__(self):
        # 类的初始化

    def __getitem__(self, item):
        # 根据索引item 返回数据

    def __len__(self):
        # 返回数据集的总数

接下来我们通过例子和代码注释来学习一下如何创建自己的数据集：

假设我们有一下文件目录结构：

filename_dataset.py 是我们创建数据集的脚本文件，我们的目的是每次读取都能返回一个shape为(batch_size, channel, seq_len)的语音数据。

import fnmatch
import os
import librosa
import numpy as np
from torch.utils.data import Dataset


# 创建fileDataset数据集
class fileDataset(Dataset):
    def __init__(self, data_folder, sr=16000, dimension=8192):
        """
        :param data_folder:音频数据地址
        :param sr: 音频采样率
        :param dimension:音频帧长
        """
        self.data_folder = data_folder
        self.sr = sr
        self.dim = dimension

        # 获取音频名列表
        self.wav_list = []
        for root, dirnames, filenames in os.walk(data_folder):
            for filename in fnmatch.filter(filenames, "*.wav"):  # 实现列表特殊字符的过滤或筛选,返回符合匹配“.wav”字符列表
                self.wav_list.append(os.path.join(root, filename))

    def __getitem__(self, item):
        # 根据索引item 返回数据
        filename = self.wav_list[item]  # 从列表中取一个音频文件名
        wb_wav, _ = librosa.load(filename, sr=self.sr)  # 读取音频文件

        # 取 帧
        if len(wb_wav) >= self.dim:
            # 如果音频长度大于帧长，则随机取一帧
            max_audio_start = len(wb_wav) - self.dim
            audio_start = np.random.randint(0, max_audio_start)
            wb_wav = wb_wav[audio_start: audio_start + self.dim]
        else:
            # 如果音频长度小于帧长，则通过在音频后面补0，补齐到帧长
            wb_wav = np.pad(wb_wav, (0, self.dim - len(wb_wav)), "constant")

        # 返回一帧语音数据和文件名
        return wb_wav, filename

    def __len__(self):
        # 返回音频文件的总数
        return len(self.wav_list)


# 实例化 fileDataset数据集 对象
train_set = fileDataset("./p225", sr=16000)

for data in train_set:
    wav_data, wav_name = data   # 第一个返回对象是数据，第二个返回对象是文件名
    print(wav_data.shape)  # (8192,)
    print(wav_name)  # 0 ./p225\p225_001.wav
    break

训练模型

搭建网络

我们先来定义一个网络，所有的使用pytorch框架定义的的神经网络模型都需要继承nn.Module类。在定义神经网络时，我们需要在__init__()函数中初始化网络层，在forward()函数将数据输入神经网络进行前向传播，返回模型输出。

import torch
import torch.nn as nn
import torch.nn.functional as F


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        # 输入图像channel:1，输出channel:6; 5*5卷积核
        self.conv1 = nn.Conv2d(1, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        # 前向传播
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        # 如果核大小是正方形，则只能指定一个数字
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = x.view(-1, self.num_flat_features(x))   # reshape 成二维，方便做全连接操作
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

    def num_flat_features(self, x):
        size = x.size()[1:]  # 除去 batch 维度的其他维度
        num_features = 1
        for s in size:
            num_features *= s
        return num_features


net = Net()     # 打印模型结构
print(net)
# Net(
#   (conv1): Conv2d(1, 6, kernel_size=(5, 5), stride=(1, 1))
#   (conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))
#   (fc1): Linear(in_features=400, out_features=120, bias=True)
#   (fc2): Linear(in_features=120, out_features=84, bias=True)
#   (fc3): Linear(in_features=84, out_features=10, bias=True))

权重初始化

对网络中的某一层进行初始化

self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3)
init.xavier_uniform(self.conv1.weight)
init.constant(self.conv1.bias, 0.1)

对网络的整体进行初始化

在网络之外初始化

方法一

def weights_init(m):
    classname = m.__class__.__name__    # 返回传入的module类型
    if classname.find("Conv") != -1:
        m.weight.data.xavier_(0.0, 0.02)
    elif classname.find("BatchNorm2d") != -1:
        m.weight.data.normal_(1.0, 0.02)    # bn层里初始化γ，服从（1，0.02）的正态分布# bn层里初始化γ，服从（1，0.02）的正态分布
        m.bias.data.fill_(0)                # bn层里初始化β，默认为0
model = Net()    # 构建网络
# 对所有的Conv层都初始化权重. 
# apply函数会递归地搜索网络内的所有module并把参数表示的函数应用到所有的module上。
model.apply(weights_init)

不建议访问以下划线为前缀的成员，他们是内部的，如果有改变不会通知用户。更推荐的一种方法是检查某个module是否是某种类型：

方法二

for m in model.modules():
    if isinstance(m, (nn.Conv2d, nn.Linear)):
        # 卷积层和全连接层参数初始化
        nn.init.normal(m.weight.data)
        m.bias.data.fill_(0)
    elif isinstance(m, nn.BatchNorm2d):
        # BatchNorm2d层参数初始化
        m.weight.data.normal_()

在网络里面初始化

def weights_init(m):
    classname = m.__class__.__name__
    if classname.find("Conv") != -1:
        m.weight.data.normal_(0.0, 0.02)
    elif classname.find("BatchNorm2d") != -1:
        m.weight.data.normal_(1.0, 0.02)
        m.bias.data.fill_(0)

def my_model(nn.Module):
    def __init__(self):
        # 因为self代表类的实例化，
        # 也就是说谁调用这个类的方法，self就指向谁
        # 我们可以在__init()__中，直接初始化模型，一般放在最后
        self.apply(weights_init)

关于pytorch中的参数初始化方法总结可以参考这篇文章：链接，

损失函数

我们这里计算均方误差 $loss=nn.MSELoss(模型预测值-目标)$

output = net(input)             # torch.Size([1, 10])
target = torch.randn(10)        # 生成一个随机数据作为target
target = target.reshape(1,-1)   # [1, 10]
mse_loss = nn.MSELoss()
loss_value = mse_loss(output, target)
print(loss_value)   # tensor(0.5513, grad_fn=<MseLossBackward>)

当我们调用loss.backward()，将误差反向传播，图中所有设置了requires_grad=True的张量的开始计算梯度，模型开始反向传播训练参数

反向传播

为了实现损失函数的梯度反向传播，我们只需要使用 loss.backward() 来反向传播权重。首先需要清零现有的梯度，否则梯度会与之前计算的梯度累加。

我们还可以观察 conv1层的偏置项(bias)的梯度，在反向传播前后的梯度。

net.zero_grad()     # 清零所有参数的梯度
print('反向传播之前的 conv1.bias.grad 梯度')
print(net.conv1.bias.grad)
# tensor([0., 0., 0., 0., 0., 0.])

loss.backward()

print('反向传播之后的 conv1.bias.grad 梯度')
print(net.conv1.bias.grad)
# tensor([-0.0118,  0.0125, -0.0085, -0.0225,  0.0125,  0.0235])

优化器更新模型参数

pytorch在 torch.optim 中集成了非常多的优化器，我们经常用的有：Adam、SGD、RMSProp...，使用方法如下

import torch.optim as optim

optimizer = optim.SGD(net.parameters(), lr=0.01)    # 创建 SGD 优化器

optimizer.zero_grad()   # 清零梯度缓存
output = net(input)
loss = criterion(output, target)    # 损失函数
loss.backward()     # 损失函数的梯度反向传播
optimizer.step()    # 更新参数

TensorBoard可视化

我通过from torch.utils.tensorboard import SummaryWriter导入tensorboard有问题，因此我选择通过tensorboardX。

from tensorboardX import SummaryWriter

创建事件对象：writer = SummaryWriter(logdir)

写入图片数据：writer.add_image(tag, img_tensor, global_step=None)

写入标量数据：writer.add_scalar(tag=, scalar_value, global_step=None)

关闭事件对象：writer.close()

在事件文件夹 ./events 中打开cmd，输入

tensorboard --logdir=runs
# 或者
tensorboard --logdir "./"

然后在浏览器中输入

https://localhost:6006/ 即可显示

MNIST图像分类器

我们通过一个小小的项目案例来讲解

导包和定义超参数
加载训练和测试数据集
定义一个卷积神经网络
定义一个损失函数
在训练样本数据上训练网络
在测试样本数据上测试网络

# Author:凌逆战
# -*- coding:utf-8 -*-
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from tensorboardX import SummaryWriter
from torchvision import datasets, transforms

batch_size = 64
epochs = 10
checkpoints_dir = "./checkpoint"
event_dir = "./enent_file"
model_name = None   # 如果要加载模型继续训练 则 "/10.pth"
lr = 1e-4

检测GPU是否可用，可用则使用GPU，不可用则使用CPU

print("GPU是否可用：", torch.cuda.is_available())  # True
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

加载MNIST数据集

# 实例化 Dataset
train_dataset = datasets.MNIST(root="./dataset/", train=True, download=True,
                               transform=transforms.Compose([
                                   transforms.ToTensor(),
                                   transforms.Normalize((0.1307,), (0.3081,))]))
test_dataset = datasets.MNIST(root="./dataset/",
                              train=False,
                              transform=transforms.Compose([
                                  transforms.ToTensor(),
                                  transforms.Normalize((0.1307,), (0.3081,))]))
# 数据加载器
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=True)

保存检查点的地址(如果检查点不存在，则创建)

# ###########    保存检查点的地址(如果检查点不存在，则创建)   ############
if not os.path.exists(checkpoints_dir):
    os.makedirs(checkpoints_dir)

模型搭建

# ###########    模型搭建   ############
class Net(nn.Module):
    """ConvNet -> Max_Pool -> RELU -> ConvNet -> Max_Pool -> RELU -> FC -> RELU -> FC -> SOFTMAX"""

    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=20, kernel_size=5, stride=1)
        self.conv2 = nn.Conv2d(in_channels=20, out_channels=50, kernel_size=5, stride=1)
        self.fc1 = nn.Linear(in_features=4 * 4 * 50, out_features=500)
        self.fc2 = nn.Linear(in_features=500, out_features=10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, kernel_size=2, stride=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 x

# 模型实例化，如果GPU可用则把模型放到GPU上
model = Net().to(device)

损失函数

# ###########    损失函数   ############
criterion = nn.CrossEntropyLoss()  # 定义交叉熵损失函数

优化器

# ###########    优化器   ############
optimizer = optim.SGD(model.parameters(), lr=lr)

TensorBoard可视化

# ###########    TensorBoard可视化 summary  ############
writer = SummaryWriter(event_dir)  # 创建事件文件

如果我们之前训练了模型并且中途中断，想要继续训练，则进行检查，如果存在则加载之前的模型继续训练

# ###########    加载模型检查点   ############
start_epoch = 0
if model_name:
    print("加载模型：", checkpoints_dir + model_name)  # "./checkpoint/10.pth"
    checkpoint = torch.load(checkpoints_dir + model_name)
    model.load_state_dict(checkpoint['model_state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    start_epoch = checkpoint['epoch']

开始训练

for epoch in range(start_epoch, epochs):
    # ###########    训练   ############
    model.train()  # 模型训练 标识
    for batch_idx, (data, target) in enumerate(train_loader):
        data = data.to(device)  # 训练数据，放到GPU上
        target = target.to(device)  # 训练标签，放到GPU上

        # 前向传播
        output = model(data)
        loss = criterion(output, target)  # 计算损失函数

        # 反向传播
        optimizer.zero_grad()  # 将梯度清零
        loss.backward()  # 反向传播
        optimizer.step()  # 更新参数

    print('Train Epoch: {} \tLoss: {:.6f}'.format(epoch+1, loss.item()))

    # ###########    TensorBoard可视化 summary  ############
    writer.add_scalar(tag="train_loss", scalar_value=loss.item(), global_step=epoch + 1)
    writer.flush()

    # ###########    测试   ############
    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)
            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 += criterion(output, target).item()

    test_loss /= len(test_loader.dataset)
    print('测试集: 损失: {:.4f}, 精度: {:.2f}%'.format(
        test_loss, 100. * correct / len(test_loader.dataset)))

    # ###########    TensorBoard可视化 summary  ############
    writer.add_scalar(tag="val_loss", scalar_value=test_loss, global_step=epoch + 1)
    writer.flush()

    # ###########    保存模型   ############
    # 每10个epoch保存一次模型
    if (epoch + 1) % 10 == 0:
        checkpoint = {
            "model_state_dict": model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            "epoch": epoch + 1,
            # 'lr_schedule': lr_schedule.state_dict()
        }
        torch.save(checkpoint, '%s/%d.pth' % (checkpoints_dir, epochs))

训练过程和结果我放在里

GPU是否可用： True
Train Epoch: 1     Loss: 2.263870
测试集: 损失: 0.0358, 精度: 12.71%
Train Epoch: 2     Loss: 2.223585
测试集: 损失: 0.0352, 精度: 36.97%
Train Epoch: 3     Loss: 2.185703
测试集: 损失: 0.0345, 精度: 49.51%
Train Epoch: 4     Loss: 2.164557
测试集: 损失: 0.0337, 精度: 57.56%
Train Epoch: 5     Loss: 2.082955
测试集: 损失: 0.0326, 精度: 63.23%
Train Epoch: 6     Loss: 1.965825
测试集: 损失: 0.0312, 精度: 66.60%
Train Epoch: 7     Loss: 1.818842
测试集: 损失: 0.0294, 精度: 68.75%
Train Epoch: 8     Loss: 1.750836
测试集: 损失: 0.0270, 精度: 71.13%
Train Epoch: 9     Loss: 1.412864
测试集: 损失: 0.0242, 精度: 74.24%
Train Epoch: 10     Loss: 1.491193
测试集: 损失: 0.0210, 精度: 77.63%

训练结果和过程

如果我们先看tensorboard的曲线，可以打开enent_file文件夹，在当前文件夹打开cmd，然后输入tensorboard --logdir "./"，就可以看到

然后在浏览器中输入https://localhost:6006/ 即可显示

多GPU训练模型

加速神经网络训练最简单的办法就是上GPU，如果一块GPU还是不够，就多上几块。像BERT和GPT-2这样的大型语言模型甚至是在上百块GPU上训练的。为了实现多GPU训练，我们必须想一个办法在多个GPU上分发数据和模型，并且协调训练过程。

单机多卡的办法有很多：

nn.DataParallel (简单方便)
torch.distributed (进阶)
apex (高级)

这里，记录了使用 4 块 Tesla V100-PICE 在 ImageNet 进行了运行时间的测试，测试结果发现 Apex 的加速效果最好，但与 Horovod/Distributed 差别不大，平时可以直接使用内置的 Distributed。Dataparallel 较慢，不推荐使用。

torch.nn.DataParallel

torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0)

DataParallel 会自动帮我们将数据切分 load 到相应 GPU，将模型复制到相应 GPU，进行正向传播计算梯度并汇总

module：要并行化的模型
device_ids：参与训练的 GPU 有哪些，(默认:所有设备)
output_device：用于汇总梯度的 GPU 是哪个，(默认:device_ids[0])

这里需要注意，模型和数据都需要先 load 进 GPU 中，DataParallel 的 module 才能对其进行处理，否则会报错：

# main.py
import os
import torch
from torch import nn
from torch.utils.data import DataLoader

os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"   # 使用第一个和第二个GPU

train_dataset = ...
train_loader = DataLoader(train_dataset, batch_size=...)

# 如果GPU可用，则环境变量CUDA_VISIBLE_DEVICES中指定的全部GPU都会被拿来使用
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 实例化模型
model = ....to(device)
if torch.cuda.device_count() > 1:
   print("Let's use", torch.cuda.device_count(), "GPUs!")
   model = nn.DataParallel(model)

optimizer = torch.optim.SGD(model.parameters(),lr=...)

for epoch in range(100):
   for batch_idx, (data, target) in enumerate(train_loader):
      images = images.to(device)
      target = target.to(device)
      ...
      output = model(images)
      loss = criterion(output, target)
      ...
      optimizer.zero_grad()
      loss.backward()
      optimizer.step()

缺点：

在每个训练批次（batch）中，因为模型的权重都是在一个进程上先算出来然后再把他们分发到每个GPU上，所以网络通信就成为了一个瓶颈，而GPU使用率也通常很低。
除此之外，nn.DataParallel 需要所有的GPU都在一个节点（一台机器）上，且并不支持 Apex 的混合精度训练。

nn.DataParallel 一个进程算权重使通信成为瓶颈，慢而且不支持混合精度训练

Torch.distributed

DataParallel 是单进程控制多 GPU，而 DistributedDataParallel 是多进程控制多 GPU，进程数等于GPU数，每个进程独享一个GPU，每个进程都会独立地执行代码。这意味着每个进程都独立地初始化模型、训练，当然，在每次迭代过程中会通过进程间通信共享梯度，整合梯度，然后独立地更新参数。通过 MPI 实现 CPU 通信，通过 NCCL 实现 GPU 通信。官方也推荐使用 DistributedDataParallel 。

分布式训练的具体流程：

1、pytorch 为我们提供了 torch.distributed.launch 启动器，用于在命令行分布式地执行 python 文件。在执行过程中，启动器会将当前进程的（其实就是 GPU的）index 通过参数传递给 python，我们可以这样获得当前进程的 index：

parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', default=-1, type=int,
                    help='node rank for distributed training')
args = parser.parse_args()
print(args.local_rank)

2、初始化进程组

在使用DDP之前，需要通过 torch.distributed.init_process_group 方法初始化进程组。这个方法会设置通信的后端（如NCCL）、进程数、主机名等参数，并建立进程之间的通信通道。

dist.init_process_group(backend='nccl',
                        init_method='tcp://127.0.0.1:12913',
                        world_size=nprocs,
                        rank=local_rank)

3、模型复制

在初始化进程组后，DDP会将模型复制到每个GPU设备上。每个GPU都会有一个独立的模型副本。

因为每个进程都会初始化一份模型，为保证模型初始化过程中生成的随机权重相同，需要设置随机种子。方法如下：

def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

4、输入分割

使用 torch.utils.data.DistributedSampler 来对数据进行分布式采样，将一个epoch里面的num_step个数据分成nprocs份，每个进程只处理其中一份数据，以确保每个进程(GPU)获取到不同的数据子集。

train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=..., sampler=train_sampler)

5、使用DDP包装模型

使用 torch.nn.parallel.DistributedDataParallel 包装模型，它能帮助我们为不同 GPU 上求得的梯度进行 all reduce（即汇总不同 GPU 计算所得的梯度，并同步计算结果）。all reduce 后不同 GPU 中模型的梯度均为 all reduce 之前各 GPU 梯度的均值：

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

6、最后，把数据和模型加载到当前进程使用的 GPU 中，正常进行正反向传播：

torch.cuda.set_device(args.local_rank)

model.cuda()

for epoch in range(100):
   for batch_idx, (data, target) in enumerate(train_loader):
      images = images.cuda(non_blocking=True)
      target = target.cuda(non_blocking=True)
      ...
      output = model(images)
      loss = criterion(output, target)
      ...
      optimizer.zero_grad()
      loss.backward()
      optimizer.step()

汇总一下，torch.distributed 并行训练部分主要与如下代码段有关：

import torch
import argparse
import torch.distributed as dist

parser = argparse.ArgumentParser()
# 负责创建 args.local_rank 变量，并接受 torch.distributed.launch 注入的值
parser.add_argument('--local_rank', default=-1, type=int,
                    help='node rank for distributed training')
args = parser.parse_args()

# 初始化分布式环境，主要用来帮助进程间通信
dist.init_process_group(backend='nccl')

# 每个进程根据自己的local_rank设置应该使用的GPU
torch.cuda.set_device(args.local_rank)
device = torch.device('cuda', args.local_rank)

# 固定随机种子
seed = 42
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)

# 分布式数据集
train_dataset = ...

train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=..., sampler=train_sampler)    # 注意这里的batch_size是每个GPU上的batch_size

# 分布式模型
model = ...
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank])

optimizer = optim.SGD(model.parameters())

# 只master进程做 logging，否则输出会很乱
if args.local_rank == 0:
    writer = SummaryWriter(comment='ddp-training')

for epoch in range(100):
   for batch_idx, (data, target) in enumerate(train_loader):
      images = images.cuda(non_blocking=True)
      target = target.cuda(non_blocking=True)
      ...
      output = model(images)
      loss = criterion(output, target)
      ...
      optimizer.zero_grad()
      loss.backward()
      optimizer.step()
      
      if args.local_rank == 0:
          writer.add_scalar(tag=, scalar_value=, global_step=)
          writer.flush()

在使用时，调用 torch.distributed.launch 启动器启动：

CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch --nproc_per_node=4 --nnodes=1 train.py

参数解释：

torch.distributed.launch：以命令行参数的方式将 args.local_rank 变量注入到每个进程中，每个进程得到的变量值都不相同。比如使用 4 个GPU的话，则 4 个进程获得的 args.local_rank 值分别为0、1、2、3
nproc_per_node：表示每个节点需要创建多少个进程(使用几个GPU就创建几个)
nnodes：表示使用几个节点，因为我们是做单机多核训练，所以设为1

详细代码参考：ddp_train.py、在 ImageNet 上的完整训练代码。

distributed 题外话

torch.distributed.barrier()

类似一个路障，进程会被拦住，直到所有进程都集合齐了才放行。

适合这样的场景：

只一个进程下载，其他进程可以使用下载好的文件；
只一个进程预处理数据，其他进程使用预处理且cache好的数据等。

模型保存与加载

模型的保存与加载与单GPU的方式有所不同。这里通通将参数以cpu的方式save进存储，因为如果是保存的GPU上参数，pth文件中会记录参数属于的GPU号，则加载时会加载到相应的GPU上，这样就会导致如果你GPU数目不够时会在加载模型时报错。

模型保存都是一致的，不过时刻记住你有多个进程在同时跑，所以会保存多个模型到存储上，如果使用共享存储就要注意文件名的问题，当然一般只在rank0进程上保存参数即可，因为所有进程的模型参数是同步的。

torch.save(model.module.cpu().state_dict(), "model.pth")

模型的加载：

param=torch.load("model.pth")

以下是huggingface/transformers代码中用到的模型保存代码

if torch.distributed.get_rank() == 0:
    model_to_save = model.module if hasattr(model, "module") else model  # Take care of distributed/parallel training
    model_to_save.save_pretrained(args.output_dir)
    tokenizer.save_pretrained(args.output_dir)

同一台机器上跑多个 ddp task

假设想在一台有4核GPU的电脑上跑两个ddp task，每个task使用两个核，很可能会需要如下错误：

RuntimeError: Address already in use
RuntimeError: NCCL error in: /opt/conda/conda-bld/pytorch_1544081127912/work/torch/lib/c10d/ProcessGroupNCCL.cpp:260, unhandled system error

原因是两个ddp task通讯地址冲突，这时候需要显示地设置每个task的地址

# 第一个task
export CUDA_VISIBLE_DEVICES="0,1" 
python -m torch.distributed.launch --nproc_per_node=2 --master_addr=127.0.0.1 --master_port=29501 train.py

# 第二个task
export CUDA_VISIBLE_DEVICES="2,3" 
python -m torch.distributed.launch --nproc_per_node=2 --master_addr=127.0.0.2 --master_port=29502 train.py

torch.multiprocessing

有的同学可能比较熟悉 torch.multiprocessing，也可以手动使用 torch.multiprocessing 进行多进程控制。绕开 torch.distributed.launch 自动控制开启和退出进程的一些小毛病～

这里有一点：需要安装NCCL，如果没有安装NCCL训练起来和单卡没有区别，我就是被这一点坑了，花了好几天时间😭。

使用时，只需要调用 torch.multiprocessing.spawn，torch.multiprocessing 就会帮助我们自动创建进程。如下面的代码所示，spawn 开启了 nprocs=4 个进程，每个进程执行 main_worker 并向其中传入 local_rank（当前进程 index）和 args（即 4 和 myargs）作为参数：

import torch.multiprocessing as mp

mp.spawn(main_worker, nprocs=4, args=(4, myargs))

这里，我们直接将原本需要 torch.distributed.launch 管理的执行内容，封装进 main_worker 函数中，其中 proc 对应 local_rank（当前进程 index），进程数 nproc 对应 4， args 对应 myargs：

def main_worker(proc, nproc, args):

   dist.init_process_group(backend='nccl', init_method='tcp://127.0.0.1:23456', world_size=4, rank=gpu)
   torch.cuda.set_device(args.local_rank)

   train_dataset = ...
   train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)

   train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=..., sampler=train_sampler)

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

   optimizer = optim.SGD(model.parameters())

   for epoch in range(100):
      for batch_idx, (data, target) in enumerate(train_loader):
          images = images.cuda(non_blocking=True)
          target = target.cuda(non_blocking=True)
          ...
          output = model(images)
          loss = criterion(output, target)
          ...
          optimizer.zero_grad()
          loss.backward()
          optimizer.step()

在上面的代码中值得注意的是，由于没有 torch.distributed.launch 读取的默认环境变量作为配置，我们需要手动为 init_process_group 指定参数：

dist.init_process_group(backend='nccl', init_method='tcp://127.0.0.1:23456', world_size=4, rank=gpu)

汇总一下，添加 multiprocessing 后并行训练部分主要与如下代码段有关：

# main.py
import torch
import torch.distributed as dist
import torch.multiprocessing as mp

mp.spawn(main_worker, nprocs=4, args=(4, myargs))

def main_worker(proc, nprocs, args):

   dist.init_process_group(backend='nccl', init_method='tcp://127.0.0.1:23456', world_size=4, rank=gpu)
   torch.cuda.set_device(args.local_rank)

   train_dataset = ...
   train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)

   train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=..., sampler=train_sampler)

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

   optimizer = optim.SGD(model.parameters())

   for epoch in range(100):
      for batch_idx, (data, target) in enumerate(train_loader):
          images = images.cuda(non_blocking=True)
          target = target.cuda(non_blocking=True)
          ...
          output = model(images)
          loss = criterion(output, target)
          ...
          optimizer.zero_grad()
          loss.backward()
          optimizer.step()

在使用时，直接使用 python 运行就可以了：

python main.py

在 ImageNet 上的完整训练代码，请点击Github。

使用Apex再加速

以后再补

Horovod 的优雅实现

Horovod 是 Uber 开源的深度学习工具，它的发展吸取了 Facebook "Training ImageNet In 1 Hour" 与百度 "Ring Allreduce" 的优点，可以无痛与 PyTorch/Tensorflow 等深度学习框架结合，实现并行训练。

在 API 层面，Horovod 和 torch.distributed 十分相似。在 mpirun 的基础上，Horovod 提供了自己封装的 horovodrun 作为启动器。

与 torch.distributed.launch 相似，我们只需要编写一份代码，horovodrun 启动器就会自动将其分配给 $n$ 个进程，分别在 $n$ 个 GPU 上运行。在执行过程中，启动器会将当前进程的（其实就是 GPU的）index 注入 hvd，我们可以这样获得当前进程的 index：

import horovod.torch as hvd

hvd.local_rank()

与 init_process_group 相似，Horovod 使用 init 设置GPU 之间通信使用的后端和端口：

hvd.init()

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

train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=..., sampler=train_sampler)

之后，使用 broadcast_parameters 包装模型参数，将模型参数从编号为 root_rank 的 GPU 复制到所有其他 GPU 中：

hvd.broadcast_parameters(model.state_dict(), root_rank=0)

然后，使用 DistributedOptimizer 包装优化器。它能帮助我们为不同 GPU 上求得的梯度进行 all reduce（即汇总不同 GPU 计算所得的梯度，并同步计算结果）。all reduce 后不同 GPU 中模型的梯度均为 all reduce 之前各 GPU 梯度的均值：

hvd.DistributedOptimizer(optimizer, named_parameters=model.named_parameters(), compression=hvd.Compression.fp16)

最后，把数据加载到当前 GPU 中。在编写代码时，我们只需要关注正常进行正向传播和反向传播：

torch.cuda.set_device(args.local_rank)

for epoch in range(100):
   for batch_idx, (data, target) in enumerate(train_loader):
      images = images.cuda(non_blocking=True)
      target = target.cuda(non_blocking=True)
      ...
      output = model(images)
      loss = criterion(output, target)
      ...
      optimizer.zero_grad()
      loss.backward()
      optimizer.step()

汇总一下，Horovod 的并行训练部分主要与如下代码段有关：

# main.py
import torch
import horovod.torch as hvd

hvd.init()
torch.cuda.set_device(hvd.local_rank())

train_dataset = ...
train_sampler = torch.utils.data.distributed.DistributedSampler(
    train_dataset, num_replicas=hvd.size(), rank=hvd.rank())

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=..., sampler=train_sampler)

model = ...
model.cuda()

optimizer = optim.SGD(model.parameters())

optimizer = hvd.DistributedOptimizer(optimizer, named_parameters=model.named_parameters())
hvd.broadcast_parameters(model.state_dict(), root_rank=0)

for epoch in range(100):
   for batch_idx, (data, target) in enumerate(train_loader):
       images = images.cuda(non_blocking=True)
       target = target.cuda(non_blocking=True)
       ...
       output = model(images)
       loss = criterion(output, target)
       ...
       optimizer.zero_grad()
       loss.backward()
       optimizer.step()

在使用时，调用 horovodrun 启动器启动：

CUDA_VISIBLE_DEVICES=0,1,2,3 horovodrun -np 4 -H localhost:4 --verbose python main.py

在 ImageNet 上的完整训练代码，请点击Github。

保存和加载模型

torch.save：保存模型，序列化对象保存到磁盘，常见的PyTorch约定是使用.pt或 .pth文件扩展名保存模型。

torch.load：加载模型，目标文件反序列化到内存中

torch.nn.Module.load_state_dict：使用反序列化的state_dict加载模型的参数字典

state_dict：python字典，包括具有可学习参数的层、每层的参数张量、优化器以及优化器超参数

为了充分了解state_dict，我们看下面例子：

import torch.nn as nn
import torch.nn.functional as F
from torch import optim


class TheModelClass(nn.Module):
    def __init__(self):
        super(TheModelClass, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


model = TheModelClass() # 初始化模型
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)   # 初始化optimizer

print("Model的state_dict:")
for param_tensor in model.state_dict():
    print(param_tensor, "\t", model.state_dict()[param_tensor].size())

print("Optimizer的state_dict:")
for var_name in optimizer.state_dict():
    print(var_name, "\t", optimizer.state_dict()[var_name])
    
# Model的state_dict:
# conv1.weight      torch.Size([6, 3, 5, 5])
# conv1.bias      torch.Size([6])
# conv2.weight      torch.Size([16, 6, 5, 5])
# conv2.bias      torch.Size([16])
# fc1.weight      torch.Size([120, 400])
# fc1.bias      torch.Size([120])
# fc2.weight      torch.Size([84, 120])
# fc2.bias      torch.Size([84])
# fc3.weight      torch.Size([10, 84])
# fc3.bias      torch.Size([10])
# Optimizer的state_dict:
# state      {}
# param_groups      [{'lr': 0.001, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'params': [3251954079208, 3251954079280, 3251954079352, 3251954079424, 3251954079496, 3251954079568, 3251954079640, 3251954079712, 3251954079784, 3251954079856]}]

保存

torch.save(model.state_dict(), PATH)  # 保存模型的参数
torch.save(model, PATH)  # 保存整个模型

加载

model.load_state_dict(torch.load(PATH))  # 加载模型的参数
model = torch.load(PATH)  # 加载整个模型

继续训练

保存

checkpoint = {
    'epoch': epoch,
    'model_state_dict': model.state_dict(), # 模型参数
    'optimizer_state_dict': optimizer.state_dict(), # 优化器参数
    'loss': loss,
    ...
    }

PATH = './checkpoint/ckpt_best_%s.pth' %(str(epoch))    # path中要包含.pth
torch.save(checkpoint, PATH)

加载

model = TheModelClass(*args, **kwargs)
optimizer = TheOptimizerClass(*args, **kwargs)

checkpoint = torch.load(PATH)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']

model.eval()
# 或
model.train()