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)
测试代码
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. 静态图
为了加速计算，一些框架会使用对神经网络“先编译，后执行”的静态图来描述网络。静态图的缺点是难以描述控制流（比如 if-else 分支语句和 for 循环语句），直接对其引入控制语句会导致产生不同的计算图。比如循环执行 n 次 a=a+b，对于不同的 n，会生成不同的计算图。
2. 推理模式
像dropout或batchnorm这样的运算符在推断和训练模式下的行为会有所不同。
# 将模型转换为推理模式
torch_model.eval()
torch_model.train(False)
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参数字典。
# 第一种：保存和加载整个模型
Save:
torch.save(model_object, 'model.pth')
Load:
model = torch.load('model.pth')
model.eval()
#第二种：仅保存和加载模型参数(推荐使用)
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. 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))
示例二
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') # 模型保存.
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")
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")
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"}})
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"}})
六、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))
七、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() 后该组参数不会变化，只可人为地改变它们的值），但是保存模型时，该组参数又作为模型参数不可或缺的一部分被保存。
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 #使用
解释：
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()中
torch.onnx.export()
模型部署入门教程（三）：PyTorch 转 ONNX 详解
torch.onnx.export中需要的模型实际上是一个torch.jit.ScriptModule。
trace跟踪法通过实际运行一遍模型的方法导出模型的静态图，即无法识别出模型中的控制流（如循环）；script记录法则能通过解析模型来正确记录所有的控制流。
scatter()
【Pytorch】scatter函数详解
target.scatter(dim, index, src)
按照指定的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]])
整个函数的操作过程见下面的示意图。因为设定了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]])
squeeze()
pytorch学习 中 torch.squeeze() 和torch.unsqueeze()的用法
功能：压缩维度，删除维度为1的维度。
torch.squeeze()，主要对数据的维度进行压缩，去掉维度为1的维度。例如，一个一行三列（1, 3）的tensor去掉第一个维数为一的维度之后，就变成了（3）行。
- squeeze(a) ，将a中所有为1的维度删掉，不为1的维度没有影响；
- a.squeeze(N)，如果第N维度为1，则删除该维度；
- b = torch.squeeze(a, N)，如果第N维度为1，则删除该维度。
unsqueeze()
功能：扩充维度，在指定位置N加上维度为1的维度。
torch.unsqueeze()，给指定位置加上维数为一的维度。例如，有个三行的tensor（3），在0的位置加了一维就变成一行三列(1, 3)。
- a.unsqueeze(N)，在a中指定位置N加上一个维数为1的维度。
- 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")
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)
