PyTorch详细教程_pytorch教程

一、参考资料

PyTorch中文文档

PyTorch官方文档

PyTorch官方源码：GitHub - pytorch/pytorch: Tensors and Dynamic neural networks in Python with strong GPU acceleration

PyTorch 中文教程 & 文档

二、分布式训练

pytorch set_epoch()方法

在分布式模式下，需要在每个 epoch 开始时调用 set_epoch() 方法，然后再创建 DataLoader迭代器，以使 shuffle 操作能够在多个 epoch 中正常工作。 否则，dataloader迭代器产生的数据将始终使用相同的顺序。

sampler = DistributedSampler(dataset) if is_distributed else None
loader = DataLoader(dataset, shuffle=(sampler is None),
                    sampler=sampler)
for epoch in range(start_epoch, n_epochs):
    if is_distributed:
        sampler.set_epoch(epoch)
    train(loader)
1
2
3
4
5
6
7

测试代码

Pytorch DistributedDataParallel 数据采样 shuffle - 知乎 (zhihu.com)

数据长度16，两张卡，每张卡8个数据，batch size 2，两个epoch，一个gpu 4次输出为一个 epoch。可以看到 cuda 1 和 cuda 0的结果是重复的。

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.distributed import DistributedSampler


torch.distributed.init_process_group(backend="nccl")

input_size = 5
output_size = 2
batch_size = 2
data_size = 16

local_rank = torch.distributed.get_rank()
torch.cuda.set_device(local_rank)
device = torch.device("cuda", local_rank)

class RandomDataset(Dataset):
    def __init__(self, size, length, local_rank):
        self.len = length
        self.data = torch.stack([torch.ones(5), torch.ones(5)*2,
                                 torch.ones(5)*3,torch.ones(5)*4,
                                 torch.ones(5)*5,torch.ones(5)*6,
                                 torch.ones(5)*7,torch.ones(5)*8,
                                 torch.ones(5)*9, torch.ones(5)*10,
                                 torch.ones(5)*11,torch.ones(5)*12,
                                 torch.ones(5)*13,torch.ones(5)*14,
                                 torch.ones(5)*15,torch.ones(5)*16]).to('cuda')

        self.local_rank = local_rank
    def __getitem__(self, index):

        return self.data[index]

    def __len__(self):
        return self.len
    
dataset = RandomDataset(input_size, data_size, local_rank)
sampler = DistributedSampler(dataset)
rand_loader = DataLoader(dataset=dataset,
                         batch_size=batch_size,
                         sampler=sampler)

e = 0
while e < 2:
    t = 0
    # 设置set_epoch()，可实现每次epoch每个GPU拿到的数据不同
    # sampler.set_epoch(e)
    for data in rand_loader:
        print(data)
    e+=1
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51

三、相关介绍

1. 静态图

为了加速计算，一些框架会使用对神经网络“先编译，后执行”的静态图来描述网络。静态图的缺点是难以描述控制流（比如 if-else 分支语句和 for 循环语句），直接对其引入控制语句会导致产生不同的计算图。比如循环执行 n 次 a=a+b，对于不同的 n，会生成不同的计算图。

2. 推理模式

像dropout或batchnorm这样的运算符在推断和训练模式下的行为会有所不同。

# 将模型转换为推理模式
torch_model.eval()
torch_model.train(False)
1
2
3

3. TorchScript格式

TorchScript

TorchScript 解读（一）：初识 TorchScript

TorchScript 是一种序列化和优化 PyTorch 模型的格式，在优化过程中，一个torch.nn.Module模型会被转换成 TorchScript 的 torch.jit.ScriptModule模型。TorchScript 也被常当成一种中间表示使用。

TorchScript 的主要用途是进行模型部署，需要记录生成一个便于推理优化的 IR，对计算图的编辑通常都是面向性能提升等等，不会给模型本身添加新的功能。

4. ATen 算子

ATen 是 PyTorch 内置的 C++ 张量计算库，PyTorch 算子在底层绝大多数计算都是用 ATen 实现的。

5. Graph

Graph 拥有许多的 Node，这些 Node 由一个 Block 管理。所有 Node 组织成双向链表的形式，方便插入删除，其中返回值节点“Return Node”会作为这个双向链表的“哨兵”。双向链表通常会被拓扑排序，保证执行的正确性。

6. Block

编译原理中，Block基本块表示一系列不包含任何跳转指令的指令序列，由于基本块内的内容可以保证是顺序执行的，因此很多的优化都会以基本块作为前提。

Block表示一个 Node 的有序列表，代表输入的 Node 的kind=Param，代表输出的 Node 的kind=Return。

实际上 Graph 本身隐含一个 root Block 对象，用来管理所有的 Node。部分 Node 可能还会存在 sub Block。

7. pass

TorchScript 解读（二）：Torch jit tracer 实现解析

pass是一个来源于编译原理的概念，一个 TorchScript 的 pass 会接收一种中间表示（IR），遍历图中所有元素进行某种变换，生成满足某种条件的新 IR。

TorchScript 中定义了许多 pass 来优化 Graph。比如对于常规编译器很常见的 DeadCodeElimination（DCE），CommonSubgraphElimination(CSE)等等；也有一些针对深度学习的融合优化，比如 FuseConvBN 等；还有针对特殊任务的 pass，ONNX 的导出就是其中一类 pass。

四、新特性

8比特优化器

github仓库

一训练就显存爆炸？Facebook 推出 8 比特优化器，两行代码拯救你的显存！

2. 超大模型训练

独家 | 如何在GPU资源受限情况下微调超大模型

五、pth模型保存与加载

Pytorch分类模型转onnx以及onnx模型推理

TorchScript 简介

1.torch.save：将序列化的对象保存到disk。这个函数使用Python的pickle实用程序进行序列化。使用这个函数可以保存各种对象的模型、张量和字典。
2.torch.load：使用pickle unpickle工具将pickle的对象文件反序列化为内存。
3.torch.nn.Module.load_state_dict:使用反序列化状态字典加载model's参数字典。
1
2
3

# 第一种：保存和加载整个模型
Save:
torch.save(model_object, 'model.pth')

Load:
model = torch.load('model.pth')
model.eval()
1
2
3
4
5
6
7

#第二种：仅保存和加载模型参数(推荐使用)
Save：
torch.save(model.state_dict(), 'params.pth')

Load:
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load('params.pth'))
model.eval()
#记住，必须调用model.eval()，以便在运行推断之前将dropout和batch规范化层设置为评估模式。如果不这样做，将会产生不一致的推断结果
#在保存用于推理或恢复训练的通用检查点时，必须保存模型的state_dict
1
2
3
4
5
6
7
8
9
10

1. trace跟踪模式

所谓 trace 指的是进行一次模型推理，在推理的过程中记录所有经过的计算，将这些记录整合成计算图，即模型的静态图。trace跟踪模式的缺点是：无法识别出模型中的控制流（如循环）。

示例一

class MyCell(torch.nn.Module):
    def __init__(self):
        super(MyCell, self).__init__()
        self.linear = torch.nn.Linear(4, 4)

    def forward(self, x, h):
        new_h = torch.tanh(self.linear(x) + h)
        return new_h, new_h

my_cell = MyCell()
x, h = torch.rand(3, 4), torch.rand(3, 4)
traced_cell = torch.jit.trace(my_cell, (x, h))
print(traced_cell)
traced_cell(x, h)

# IR中间表示
print(traced_cell.graph)
print(traced_cell.code)

# 调用traced_cell会产生与 Python 模块相同的结果
print(my_cell(x, h))
print(traced_cell(x, h))
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

示例二

class MyModule(nn.Module):
    def __init__(self):
        super(MyModule,self).__init__()
        self.conv1 = nn.Conv2d(1,3,3)

    def forward(self,x):
    	x = self.conv1(x)
        return x

model = MyModule()  # 实例化模型
trace_module = torch.jit.trace(model,torch.rand(1,1,224,224)) 
print(trace_module.code)  # 查看模型结构
output = trace_module (torch.ones(1, 3, 224, 224)) # 测试
print(output)
trace_modult('model.pt') # 模型保存.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

2. script记录模式（带控制流）

script记录模式，通过解析模型来正确记录所有的控制流。script记录模式直接解析网络定义的 python 代码，生成抽象语法树 AST。

import torch
import os
from pointnet2_cls_won_model import *
# 用于定义存储的路径
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
ROOT_DIR = BASE_DIR
# 加载权重文件
checkpoint = torch.load(ROOT_DIR + '/checkpoints/best_model.pth')
# 获取模型
classifier = get_model(num_class=10, normal_channel=False)
# 将参数加载到模型中
classifier.load_state_dict(checkpoint['model_state_dict'])
# 输出为.pt文件
scripted_gate = torch.jit.script(classifier)
scripted_gate.save(ROOT_DIR + "script_model_1.pt")
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

3. trace模型转换

import torch
import torchvision
from unet import UNet
model = UNet(3, 2)#自己定义的网络模型
model.load_state_dict(torch.load("best_weights.pth"))#保存的训练模型
model.eval()#切换到eval（）
example = torch.rand(1, 3, 320, 480)#生成一个随机输入维度的输入
traced_script_module = torch.jit.trace(model, example)
traced_script_module.save("model.pt")
1
2
3
4
5
6
7
8
9

4. 如果保存的是整个模型

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

model = torch.load("test.pth") # pytorch模型加载
batch_size = 1  #批处理大小
input_shape = (3, 244, 384)   #输入数据,改成自己的输入shape

# #set the model to inference mode
model.eval()

x = torch.randn(batch_size, *input_shape)   # 生成张量
x = x.to(device)
export_onnx_file = "test.onnx"		# 目的ONNX文件名
torch.onnx.export(model
                    x,
                    export_onnx_file,
                    opset_version=10,
                    do_constant_folding=True,	# 是否执行常量折叠优化
                    input_names=["input"],	# 输入名
                    output_names=["output"],	# 输出名
                    dynamic_axes={"input":{0:"batch_size"},  # 批处理变量
                                    "output":{0:"batch_size"}})
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

5. 如果保存的是模型参数

import torch
import torchvision.models as models

torch_model = torch.load("test.pth") # pytorch模型加载

model = models.resnet50()
model.fc = torch.nn.Linear(2048, 4)
model.load_state_dict(torch_model) 

batch_size = 1  #批处理大小
input_shape = (3, 244, 384)   #输入数据,改成自己的输入shape

# #set the model to inference mode
model.eval()

x = torch.randn(batch_size, *input_shape)	# 生成张量
export_onnx_file = "test.onnx"			# 目的ONNX文件名
torch.onnx.export(model,
                    x,
                    export_onnx_file,
                    opset_version=10,
                    do_constant_folding=True,	# 是否执行常量折叠优化
                    input_names=["input"],	# 输入名
                    output_names=["output"],	# 输出名
                    dynamic_axes={"input":{0:"batch_size"},  # 批处理变量
                                    "output":{0:"batch_size"}})
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

六、Torch工具与组件

1. TorchServe

TorchServe官方仓库

TorchServe 详解：5 步将模型部署到生产环境

torchserve使用教程（踩坑记录）

TorchServe 是 PyTorch 中将模型部署到生产环境的首选解决方案。它是一个性能良好且可扩展的工具，用 HTTP 或 HTTPS API 封装模型。

2. torchstat

torchstat
Flops counting tool for neural networks in pytorch framework

计算pytorch构建的网络的parameters，FLOPs，MAdd，内存空间大小等指标，简单好用。

This tools can show

• Total number of network parameters
• Theoretical amount of floating point arithmetics (FLOPs)
• Theoretical amount of multiply-adds (MAdd)
• Memory usage

示例

统计alexnet相关参数。

from torchstat import stat
import torchvision.models as models
 
model = model.alexnet()
stat(model, (3, 224, 224))
1
2
3
4
5

七、ONNX模型相关

1. ONNX模型推理

2. ONNX模型转换

省去编译转换工具的时间 开箱即用，一键转换

模型部署入门教程（三）：PyTorch 转 ONNX 详解

torch.onnx.export

八、Pytorch模型迁移学习

Pytorch模型迁移和迁移学习,导入部分模型参数

1. libtorch模型部署

三维目标检测：（五）如何将pytorch模型部署到C++工程中及pytorch模型转libtorch模型常见的问题

九、常用API

torch.nn.Module.register_buffer()

PyTorch nn.Module中的self.register_buffer()解析

pytorch 中register_buffer（）

深入理解Pytorch之register_buffer

功能：训练不更新，最后可保存。定义一组参数，该组参数的特别之处在于，模型训练时不会更新（即调用 optimizer.step() 后该组参数不会变化，只可人为地改变它们的值），但是保存模型时，该组参数又作为模型参数不可或缺的一部分被保存。
1

import torch.nn as nn
import torch
class net(nn.Module):
    def __init__(self):
        super(net,self).__init__()
        self.register_buffer("a",torch.ones(2,3))  #从此，self.a其实就是torch.ones(2,3)。
    def forward(self,x):
        return x+self.a  #使用
1
2
3
4
5
6
7
8

解释：

register_buffer的作用是将 torch.ones(2,3) 这个tensor注册到模型的 buffers() 属性中，并命名为a，这代表a对应的是一个 持久态，不会梯度更新，但是能被模型的state_dict记录下来。可以理解为模型的常数。

注意，没有保存到模型的 buffers() 或 parameters() 属性中的参数是不会被记录到state_dict中的，在 buffers() 中的参数默认不会有梯度，parameters() 中的参数默认有梯度。

requires_grad=False

不会注册到模型参数中model.parameters()
会注册到模型model.state_dict()中
1
2
3
4

torch.onnx.export()

模型部署入门教程（三）：PyTorch 转 ONNX 详解

torch.onnx.export中需要的模型实际上是一个torch.jit.ScriptModule。

trace跟踪法通过实际运行一遍模型的方法导出模型的静态图，即无法识别出模型中的控制流（如循环）；script记录法则能通过解析模型来正确记录所有的控制流。

scatter()

【Pytorch】scatter函数详解

target.scatter(dim, index, src)
1

按照指定的dim轴方向和index对应位置关系，用src张量中的元素逐个映射到target张量中的元素。

参数解释

• target：即目标张量；
• src：即源张量，将该张量上的元素逐个映射到目标张量上；
• dim：指定轴方向，定义了填充方式。对于二维张量，dim=0表示逐列进行行填充，而dim=1表示逐列进行行填充；
• index: 按照轴方向，在target张量中需要填充的位置；

import torch
a = torch.arange(10).reshape(2,5).float()
b = torch.zeros(3, 5))
b_= b.scatter(dim=0, index=torch.LongTensor([[1, 2, 1, 1, 2], [2, 0, 2, 1, 0]]),src=a)
print(b_)

# tensor([[0, 6, 0, 0, 9],
#        [0, 0, 2, 8, 0],
#        [5, 1, 7, 0, 4]])
1
2
3
4
5
6
7
8
9

整个函数的操作过程见下面的示意图。因为设定了dim=0，所以会逐列将source中的元素按照index中的位置信息，放入target张量中。

scatter函数的一个典型应用就是在分类问题中，将目标标签转换为one-hot编码形式，如：

labels = torch.LongTensor([1,3])
targets = torch.zeros(2, 5)
targets.scatter(dim=1, index=labels.unsqueeze(-1), src=torch.tensor(1))
# 注意dim=1，即逐样本的进行列填充
# 返回值为 tensor([[0, 1, 0, 0, 0],
#        [0, 0, 0, 1, 0]])
1
2
3
4
5
6

squeeze()

pytorch学习 中 torch.squeeze() 和torch.unsqueeze()的用法

功能：压缩维度，删除维度为1的维度。

torch.squeeze()，主要对数据的维度进行压缩，去掉维度为1的维度。例如，一个一行三列（1, 3）的tensor去掉第一个维数为一的维度之后，就变成了（3）行。

1. squeeze(a) ，将a中所有为1的维度删掉，不为1的维度没有影响；
2. a.squeeze(N)，如果第N维度为1，则删除该维度；
3. b = torch.squeeze(a, N)，如果第N维度为1，则删除该维度。

unsqueeze()

功能：扩充维度，在指定位置N加上维度为1的维度。

torch.unsqueeze()，给指定位置加上维数为一的维度。例如，有个三行的tensor（3），在0的位置加了一维就变成一行三列(1, 3)。

1. a.unsqueeze(N)，在a中指定位置N加上一个维数为1的维度。
2. b = torch.unsqueeze(a, N)，在a中指定位置N加上一个维度为1的维度。

torch.cuda

# Use the GPU if there is one, otherwise CPU
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
1
2
3
4

torch.onnx.export()

十、示例代码

torch.onnx.export()

模型部署入门教程（三）：PyTorch 转 ONNX 详解

import torch 
 
class Model(torch.nn.Module): 
    def __init__(self, n): 
        super().__init__() 
        self.n = n 
        self.conv = torch.nn.Conv2d(3, 3, 3) 
 
    def forward(self, x): 
        for i in range(self.n): 
            x = self.conv(x) 
        return x 
 
 
models = [Model(2), Model(3)] 
model_names = ['model_2', 'model_3'] 
 
for model, model_name in zip(models, model_names): 
    dummy_input = torch.rand(1, 3, 10, 10) 
    dummy_output = model(dummy_input) 
    model_trace = torch.jit.trace(model, dummy_input) 
    model_script = torch.jit.script(model) 
 
    # 跟踪法与直接 torch.onnx.export(model, ...)等价 
    torch.onnx.export(model_trace, dummy_input, f'{model_name}_trace.onnx', example_outputs=dummy_output) 

    # 记录法必须先调用 torch.jit.script 
    torch.onnx.export(model_script, dummy_input, f'{model_name}_script.onnx', example_outputs=dummy_output) 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28