- 1.Pytorch中的view、reshape方法的异同
- 2.PyTorch矩阵乘法详解
- 3.PyTorch维度变化操作详解
- 4.PyTorch模型构建详解
- 5.PyTorch中的Module
- 6.PyTorch中常用的随机采样
- 7.PyTorch中对梯度计算的控制
- 8.Pytorch中的多卡训练
- 9.DeepSpeed介绍
- 10.PyTorch中的模块迭代器
- 11.PyTorch中的DataLoader介绍
- 12.PyTorch中的动态图和静态图介绍
- 13.PyTorch中的compile介绍
- 14.AI模型训练过程的可视化实用工具有哪些?
- 15.PyTorch2.0组件TorchDynamo介绍
- 16.PyTorch2.0组件AOTAutograd介绍
- 17.介绍一下PyTorch中.detach()、.clone()、requires_grad=True、torch.no_grad()的原理与作用
- 18.PyTorch中连续张量和非连续张量有哪些区别?
- 19.OpenAI-Triton介绍
- 20.Triton实现add算子
- 21.Triton常用API介绍
- 22.Triton的编译流程
- 23.什么是IR表示?
- 24.计算图优化的常用方法?
- 25.介绍一下AI模型中钩子函数的原理和作用
- 26.介绍一下PyTorch中DataLoader库的底层原理
- 27.AIGC时代主流的深度学习框架有哪些?各自有什么特点?
- 28.Caffe的depthwise为什么慢,该如何优化?Caffe新加一层需要哪些操作?
- 29.介绍一下torchrun的用法
要想深入理解view和reshape方法的区别,我们需要先知道Pytorch中的Tensor是如何储存的。
Pytorch中tensor采用分开储存的形式,分为头信息区(Tensor)和存储区(Storage)。tensor的形状(size)、步长(stride)、数据类型(type)等信息储存在头部信息区,而真正的数据则存储在存储区。
举个例子
import torch
a = torch.arange(5) # 初始化张量 a 为 [0, 1, 2, 3, 4]
b = a[2:] # 截取张量a的部分值并赋值给b,b其实只是改变了a对数据的索引方式
print('a:', a)
print('b:', b)
print('ptr of storage of a:', a.storage().data_ptr()) # 打印a的存储区地址
print('ptr of storage of b:', b.storage().data_ptr()) # 打印b的存储区地址,可以发现两者是共用存储区
print('==================================================================')
b[1] = 0 # 修改b中索引为1,即a中索引为3的数据为0
print('a:', a)
print('b:', b)
print('ptr of storage of a:', a.storage().data_ptr()) # 打印a的存储区地址
print('ptr of storage of b:', b.storage().data_ptr()) # 打印b的存储区地址,可以发现两者是共用存储区
''' 运行结果 '''
a: tensor([0, 1, 2, 3, 4])
b: tensor([2, 3, 4])
ptr of storage of a: 2862826251264
ptr of storage of b: 2862826251264
==================================================================
a: tensor([0, 1, 2, 0, 4])
b: tensor([2, 0, 4])
ptr of storage of a: 2862826251264
ptr of storage of b: 2862826251264以发现a、b这两个tensor的Storage都是一样的,但它们的头信息区不同。
官方文档描述:stride是在指定维度dim中从一个元素跳到下一个元素所必需的步长。 举个例子
import torch
x = torch.tensor([[1, 3, 5, 7], [7, 7, 7, 7]])
print(x)
print(x.stride(0)) # 打印第0维度中第一个元素到下一个元素的步长
print(x.stride(1)) # 打印第1维度中第一个元素到下一个元素的步长
''' 运行结果 '''
tensor([[1, 3, 5, 7],
[7, 7, 7, 7]])
4
1view方法能够将tensor转换为指定的shape,且原始的data不改变。返回的tensor与原始的tensor共享存储区。但view方法需要满足以下连续条件:
举个例子,我们初始化一个tensor a与b
import torch
a = torch.arange(9).reshape(3, 3) # 初始化张量a
b = a.permute(1, 0) # 令b等于a的转置
print(a) # 打印a
print(a.size()) # 查看a的shape
print(a.stride()) # 查看a的stride
print('==================================================================')
print(b) # 打印b
print(b.size()) # 查看b的shape
print(b.stride()) # 查看b的stride
''' 运行结果 '''
tensor([[0, 1, 2],
[3, 4, 5],
[6, 7, 8]])
torch.Size([3, 3])
(3, 1)
==================================================================
tensor([[0, 3, 6],
[1, 4, 7],
[2, 5, 8]])
torch.Size([3, 3])
(1, 3)我们将tensor a与b分别带入连续性条件公式进行验证,发现a可以满足而b不满足,下面我们尝试对tensor a与b进行view操作
import torch
a = torch.arange(9).reshape(3, 3) # 初始化张量a
b = a.permute(1, 0) # 令b等于a的转置
print(a.view(-1))
''' 运行结果 '''
tensor([0, 1, 2, 3, 4, 5, 6, 7, 8])import torch
a = torch.arange(9).reshape(3, 3) # 初始化张量a
b = a.permute(1, 0) # 令b等于a的转置
print(b.view(-1))
''' 运行结果 '''
Traceback (most recent call last):
File "C:/Users/97987/PycharmProjects/pytorch/test.py", line 4, in <module>
print(b.view(-1))
RuntimeError: view size is not compatible with input tensor's size and stride (at least one dimension spans across two contiguous subspaces). Use .reshape(...) instead.果然只有在满足连续性条件下才可以使用view方法。 如果不满足此条件,则需要先使用contiguous方法将原始tensor转换为满足连续条件的tensor,然后再使用view方法进行shape变换。但是经过contiguous方法变换后的tensor将重新开辟一个储存空间,不再与原始tensor共享内存。
import torch
a = torch.arange(9).reshape(3, 3) # 初始化张量a
b = a.permute(1, 0) # 令b等于a的转置
c = b.contiguous() # 使用contiguous方法
print(c.view(-1))
print(a.storage().data_ptr())
print(b.storage().data_ptr())
print(c.storage().data_ptr())
''' 运行结果 '''
tensor([0, 3, 6, 1, 4, 7, 2, 5, 8])
2610092185792
2610092185792
2610092184704从以上结果可以看到,tensor a与c是属于不同存储区的张量,也就是说经过contiguous方法变换后的tensor将重新开辟一个储存空间,不再与原始tensor共享内存。
与view方法类似,将输入tensor转换为新的shape格式,但是reshape方法是view方法与contiguous方法的综合。 也就是说当tensor满足连续性条件时,reshape方法返回的结果与view方法相同,否则返回的结果与先经过contiguous方法在进行view方法的结果相同。
view方法和reshape方法都可以用来更改tensor的shape,但view只适合对满足连续性条件的tensor进行操作,而reshape同时还可以对不满足连续性条件的tensor进行操作,兼容性更好,而view方法可以节省内存,如果不满足连续性条件使用reshape方法则会重新开辟储存空间。
PyTorch作为深度学习领域的主流框架之一,提供了多种矩阵乘法操作。本文将详细介绍PyTorch中的各种矩阵乘法函数,帮助您在不同场景下选择最适合的方法。
torch.matmul()是PyTorch中最通用的矩阵乘法函数,可以处理多维张量。
- 支持广播机制
- 可以处理1维到4维的张量
- 根据输入张量的维度自动选择适当的乘法操作
import torch
a = torch.randn(2, 3)
b = torch.randn(3, 4)
c = torch.matmul(a, b) # 结果形状为 (2, 4)
# 也可以用@运算符
c = a @ btorch.mm()专门用于2维矩阵相乘。
- 只能处理2维矩阵
- 比
torch.matmul()在某些情况下更快
a = torch.randn(2, 3)
b = torch.randn(3, 4)
c = torch.mm(a, b) # 结果形状为 (2, 4)torch.bmm()用于批量矩阵乘法,处理3维张量。
- 输入必须是3维张量
- 用于同时计算多个矩阵乘法
a = torch.randn(10, 3, 4)
b = torch.randn(10, 4, 5)
c = torch.bmm(a, b) # 结果形状为 (10, 3, 5)Python 3.5+引入的矩阵乘法运算符,在PyTorch中也可使用。
- 语法简洁
- 功能等同于
torch.matmul()
a = torch.randn(2, 3)
b = torch.randn(3, 4)
c = a @ b # 结果形状为 (2, 4)torch.dot()计算两个一维张量的点积。
- 只能用于1维张量
- 返回一个标量
a = torch.randn(5)
b = torch.randn(5)
c = torch.dot(a, b) # 结果是一个标量torch.mv()用于矩阵与向量相乘。
- 第一个参数必须是2维矩阵
- 第二个参数必须是1维向量
matrix = torch.randn(3, 4)
vector = torch.randn(4)
result = torch.mv(matrix, vector) # 结果形状为 (3,)torch.einsum()使用爱因斯坦求和约定,可以执行更复杂的张量运算,包括矩阵乘法。
- 非常灵活,可以表达复杂的张量运算
- 语法简洁但可能难以理解
a = torch.randn(2, 3)
b = torch.randn(3, 4)
c = torch.einsum('ij,jk->ik', a, b) # 等同于矩阵乘法,结果形状为 (2, 4)PyTorch提供了多种矩阵乘法操作,适用于不同的场景:
- 对于一般情况,使用
torch.matmul()或@运算符 - 对于2维矩阵乘法,可以使用
torch.mm() - 对于批量矩阵乘法,使用
torch.bmm() - 对于向量点积,使用
torch.dot() - 对于矩阵与向量相乘,使用
torch.mv() - 对于更复杂的张量运算,可以考虑
torch.einsum()
选择合适的函数可以提高代码的可读性和运行效率。在实际应用中,建议根据具体情况选择最合适的方法。
PyTorch作为深度学习领域的主流框架,提供了丰富的维度变化操作。这些操作在数据预处理、模型构建和结果处理中都扮演着重要角色。本文将详细介绍PyTorch中的各种维度变化操作,帮助您更好地理解和使用这些功能。
这两个函数用于改变张量的形状,但不改变其数据。
- 要求张量在内存中是连续的
- 不会复制数据,只是改变视图
import torch
x = torch.randn(4, 4)
y = x.view(16)
z = x.view(-1, 8) # -1表示这个维度的大小将被自动计算- 类似于view(),但可以处理非连续的张量
- 如果可能,不会复制数据
a = torch.randn(4, 4)
b = a.reshape(2, 8)这对函数用于移除或添加维度。
移除大小为1的维度
x = torch.zeros(2, 1, 3, 1, 4)
y = x.squeeze() # y.shape: (2, 3, 4)
z = x.squeeze(1) # z.shape: (2, 3, 1, 4)在指定位置添加大小为1的维度
x = torch.tensor([1, 2, 3])
y = x.unsqueeze(0) # y.shape: (1, 3)
z = x.unsqueeze(1) # z.shape: (3, 1)这两个函数用于交换维度。
交换两个指定的维度
x = torch.randn(2, 3, 5)
y = x.transpose(0, 2) # y.shape: (5, 3, 2)可以对任意维度进行重新排列
x = torch.randn(2, 3, 5)
y = x.permute(2, 0, 1) # y.shape: (5, 2, 3)这两个函数用于扩展tensor的大小。
- 不会分配新内存,只是创建一个新的视图
- 只能扩展大小为1的维度
x = torch.tensor([[1], [2], [3]])
y = x.expand(3, 4) # y: [[1, 1, 1, 1], [2, 2, 2, 2], [3, 3, 3, 3]]- 会分配新内存,复制数据
- 可以沿着任意维度重复tensor
x = torch.tensor([1, 2, 3])
y = x.repeat(2, 3) # y: [[1, 2, 3, 1, 2, 3, 1, 2, 3], [1, 2, 3, 1, 2, 3, 1, 2, 3]]这两个函数用于将多维张量展平成一维。
将张量展平成一维
x = torch.randn(2, 3, 4)
y = x.flatten() # y.shape: (24,)
z = x.flatten(start_dim=1) # z.shape: (2, 12)功能类似于flatten(),但返回的可能是一个视图
x = torch.randn(2, 3, 4)
y = x.ravel() # y.shape: (24,)这两个函数用于连接张量。
沿着新维度连接张量
x = torch.randn(3, 4)
y = torch.randn(3, 4)
z = torch.stack([x, y]) # z.shape: (2, 3, 4)沿着已存在的维度连接张量
x = torch.randn(2, 3)
y = torch.randn(2, 5)
z = torch.cat([x, y], dim=1) # z.shape: (2, 8)这两个函数用于将张量分割成多个部分。
将张量分割成指定大小的块
x = torch.randn(5, 10)
y = torch.split(x, 2, dim=0) # 返回一个元组,包含3个tensor,形状分别为(2, 10), (2, 10), (1, 10)将张量均匀分割成指定数量的块
x = torch.randn(5, 10)
y = torch.chunk(x, 3, dim=1) # 返回一个元组,包含3个tensor,形状分别为(5, 4), (5, 3), (5, 3)将张量广播到指定的形状。
x = torch.randn(3, 1)
y = torch.broadcast_to(x, (3, 5)) # y.shape: (3, 5)可以用来缩小张量的某个维度。
x = torch.randn(3, 5)
y = x.narrow(1, 1, 2) # 在第1维(列)上,从索引1开始,选择2个元素将张量的某个维度展开。
x = torch.arange(1, 8)
y = x.unfold(0, 3, 1) # 步长为1的滑动窗口操作,窗口大小为3PyTorch提供了丰富的维度变化操作,可以满足各种数据处理和模型构建的需求:
- 改变形状:view(), reshape()
- 添加/删除维度:squeeze(), unsqueeze()
- 交换维度:transpose(), permute()
- 扩展大小:expand(), repeat()
- 展平:flatten(), ravel()
- 连接:stack(), cat()
- 分割:split(), chunk()
- 广播:broadcast_to()
- 裁剪和展开:narrow(), unfold()
熟练掌握这些操作可以帮助你更高效地处理张量数据,构建复杂的神经网络模型。
PyTorch是一个强大的深度学习框架,提供了丰富的工具和组件用于构建各种类型的神经网络模型。本文将全面介绍PyTorch中用于模型构建的主要操作和组件。
nn.Module是PyTorch中所有神经网络模块的基类。自定义模型通常继承自这个类。
import torch.nn as nn
class MyModel(nn.Module):
def __init__(self):
super(MyModel, self).__init__()
self.layer1 = nn.Linear(10, 20)
self.layer2 = nn.Linear(20, 2)
def forward(self, x):
x = torch.relu(self.layer1(x))
return self.layer2(x)
model = MyModel()nn.Sequential是一个有序的模块容器,用于快速构建线性结构的网络。
model = nn.Sequential(
nn.Linear(10, 20),
nn.ReLU(),
nn.Linear(20, 2)
)linear_layer = nn.Linear(in_features=10, out_features=20)conv_layer = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, stride=1, padding=1)rnn_layer = nn.RNN(input_size=10, hidden_size=20, num_layers=2)
lstm_layer = nn.LSTM(input_size=10, hidden_size=20, num_layers=2)
gru_layer = nn.GRU(input_size=10, hidden_size=20, num_layers=2)transformer_layer = nn.Transformer(d_model=512, nhead=8, num_encoder_layers=6, num_decoder_layers=6)PyTorch提供了多种激活函数:
relu = nn.ReLU()
sigmoid = nn.Sigmoid()
tanh = nn.Tanh()
leaky_relu = nn.LeakyReLU(negative_slope=0.01)常用的池化层包括最大池化和平均池化:
max_pool = nn.MaxPool2d(kernel_size=2, stride=2)
avg_pool = nn.AvgPool2d(kernel_size=2, stride=2)归一化层有助于稳定训练过程:
batch_norm = nn.BatchNorm2d(num_features=16)
layer_norm = nn.LayerNorm(normalized_shape=[20, 30])PyTorch提供了多种损失函数:
mse_loss = nn.MSELoss()
cross_entropy_loss = nn.CrossEntropyLoss()
bce_loss = nn.BCELoss()优化器用于更新模型参数:
import torch.optim as optim
model = MyModel()
sgd_optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
adam_optimizer = optim.Adam(model.parameters(), lr=0.001)正确的参数初始化对模型训练很重要:
def init_weights(m):
if isinstance(m, nn.Linear):
nn.init.xavier_uniform_(m.weight)
nn.init.zeros_(m.bias)
model = MyModel()
model.apply(init_weights)保存和加载模型是很常见的操作:
# 保存模型
torch.save(model.state_dict(), 'model.pth')
# 加载模型
model = MyModel()
model.load_state_dict(torch.load('model.pth'))
model.eval()对于多GPU训练,可以使用DataParallel:
model = nn.DataParallel(model)可以通过继承nn.Module来创建自定义层:
class MyCustomLayer(nn.Module):
def __init__(self, in_features, out_features):
super(MyCustomLayer, self).__init__()
self.linear = nn.Linear(in_features, out_features)
def forward(self, x):
return torch.sigmoid(self.linear(x))
custom_layer = MyCustomLayer(10, 5)这里是一个基本的训练循环示例:
model = MyModel()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
for epoch in range(num_epochs):
for inputs, labels in dataloader:
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()在训练后评估模型性能:
model.eval()
with torch.no_grad():
for inputs, labels in test_dataloader:
outputs = model(inputs)
# 计算准确率或其他指标PyTorch提供了丰富的工具和组件用于构建各种类型的神经网络模型。从基本的层和激活函数,到高级的优化器和并行处理,PyTorch都提供了强大的支持。熟练掌握这些组件和操作可以帮助你更高效地设计和实现复杂的深度学习模型。
PyTorch 使用模块(modules)来表示神经网络。模块具有以下特性:
-
构建状态计算的基石。PyTorch 提供了一个强大的模块库,并且简化了定义新自定义模块的过程,从而轻松构建复杂的多层神经网络。
-
与 PyTorch 的自动微分系统紧密集成。模块使得指定 PyTorch 优化器要更新的可学习参数变得简单。
-
易于操作和转换。模块可以方便地保存和恢复,且可以在 CPU / GPU / TPU 设备之间转换、剪枝、量化等。
首先,让我们看一个简单的自定义版本的 PyTorch 的 Linear 模块。这个模块对其输入应用仿射变换。
import torch
from torch import nn
class MyLinear(nn.Module):
def __init__(self, in_features, out_features):
super().__init__()
self.weight = nn.Parameter(torch.randn(in_features, out_features))
self.bias = nn.Parameter(torch.randn(out_features))
def forward(self, input):
return (input @ self.weight) + self.bias这个简单模块具备了模块的基本特征:
-
继承自
Module基类。所有模块应该继承自Module以便与其他模块组合。 -
定义了一些在计算中使用的“状态”。这里,状态由随机初始化的权重和偏置张量组成,这些张量定义了仿射变换。因为每个都是
Parameter,所以它们会自动注册为模块的参数,并且在调用parameters()时会返回。参数可以看作是模块计算的“可学习”方面(更多内容在后面)。请注意,模块不是必须有状态的,也可以是无状态的。 -
定义了一个执行计算的
forward()函数。对于这个仿射变换模块,输入与权重参数进行矩阵相乘(使用@符号)并加上偏置参数以生成输出。更一般地,模块的forward()实现可以执行涉及任意数量输入和输出的任意计算。
这个简单模块演示了模块如何将状态和计算打包在一起。可以构建和调用此模块的实例:
m = MyLinear(4, 3)
sample_input = torch.randn(4)
m(sample_input)
# tensor([-0.3037, -1.0413, -4.2057], grad_fn=<AddBackward0>)注意模块本身是可调用的,调用它会触发其 forward() 函数。这个名字是参考“前向传递”和“反向传递”的概念,适用于每个模块。前向传递负责将模块表示的计算应用于给定输入(如上所示)。反向传递计算模块输出相对于其输入的梯度,可以用于通过梯度下降方法“训练”参数。PyTorch 的自动微分系统会自动处理这个反向传递计算,因此不需要为每个模块手动实现 backward() 函数。通过连续的前向/反向传递来训练模块参数的过程将在“使用模块进行神经网络训练”一节中详细介绍。
可以通过调用 parameters() 或 named_parameters() 迭代模块注册的所有参数,后者包括每个参数的名称:
for parameter in m.named_parameters():
print(parameter)
# ('weight', Parameter containing:
# tensor([[ 1.0597, 1.1796, 0.8247],
# [-0.5080, -1.2635, -1.1045],
# [ 0.0593, 0.2469, -1.4299],
# [-0.4926, -0.5457, 0.4793]], requires_grad=True))
# ('bias', Parameter containing:
# tensor([ 0.3634, 0.2015, -0.8525], requires_grad=True))通常,模块注册的参数是模块计算中应该“学习”的方面。本文后面的部分将展示如何使用 PyTorch 的优化器更新这些参数。在此之前,让我们先看看模块如何相互组合。
模块可以包含其他模块,使其成为开发更复杂功能的有用构建块。最简单的方法是使用 Sequential 模块。它允许我们将多个模块串联在一起:
net = nn.Sequential(
MyLinear(4, 3),
nn.ReLU(),
MyLinear(3, 1)
)
sample_input = torch.randn(4)
net(sample_input)
# tensor([-0.6749], grad_fn=<AddBackward0>)注意 Sequential 自动将第一个 MyLinear 模块的输出作为输入传递给 ReLU,然后将其输出作为输入传递给第二个 MyLinear 模块。如所示,它仅限于具有单一输入和输出的模块的按顺序链接。
一般来说,建议为简单用例之外的任何情况定义自定义模块,因为这提供了对子模块用于模块计算的完全灵活性。
例如,下面是一个简单神经网络实现为自定义模块:
import torch.nn.functional as F
class Net(nn.Module):
def __init__(self):
super().__init__()
self.l0 = MyLinear(4, 3)
self.l1 = MyLinear(3, 1)
def forward(self, x):
x = self.l0(x)
x = F.relu(x)
x = self.l1(x)
return x该模块由定义神经网络层的两个“子模块”(l0 和 l1)组成,并在模块的 forward() 方法中用于计算。可以通过调用 children() 或 named_children() 迭代模块的直接子模块:
net = Net()
for child in net.named_children():
print(child)
# ('l0', MyLinear())
# ('l1', MyLinear())要深入到直接子模块,可以递归调用 modules() 和 named_modules() 迭代一个模块及其子模块:
class BigNet(nn.Module):
def __init__(self):
super().__init__()
self.l1 = MyLinear(5, 4)
self.net = Net()
def forward(self, x):
return self.net(self.l1(x))
big_net = BigNet()
for module in big_net.named_modules():
print(module)
# ('', BigNet(
# (l1): MyLinear()
# (net): Net(
# (l0): MyLinear()
# (l1): MyLinear()
# )
# ))
# ('l1', MyLinear())
# ('net', Net(
# (l0): MyLinear()
# (l1): MyLinear()
# ))
# ('net.l0', MyLinear())
# ('net.l1', MyLinear())有时,模块需要动态定义子模块。这时 ModuleList 和 ModuleDict 模块很有用,它们从列表或字典中注册子模块:
class DynamicNet(nn.Module):
def __init__(self, num_layers):
super().__init__()
self.linears = nn.ModuleList(
[MyLinear(4, 4) for _ in range(num_layers)])
self.activations = nn.ModuleDict({
'relu': nn.ReLU(),
'lrelu': nn.LeakyReLU()
})
self.final = MyLinear(4, 1)
def forward(self, x, act):
for linear in self.linears:
x = linear(x)
x = self.activations[act](x)
x = self.final(x)
return x
dynamic_net = DynamicNet(3)
sample_input = torch.randn(4)
output = dynamic_net(sample_input, 'relu')对于任何给定的模块,它的参数包括其直接参数以及所有子模块的参数。这意味着调用 parameters() 和 named_parameters() 会递归包含子参数,从而方便地优化网络内的所有参数:
for parameter in dynamic_net.named_parameters():
print(parameter)
# ('linears.0.weight', Parameter containing:
# tensor([[-1.2051, 0.7601, 1.1065, 0.1963],
# [ 3.0592, 0.4354, 1.6598, 0.9828],
# [-0.4446, 0.4628, 0.8774, 1.6848],
# [-0.1222, 1.5458, 1.1729, 1.4647]], requires_grad=True))
# ('linears.0.bias', Parameter containing:
# tensor([ 1.5310, 1.0609, -2.0940, 1.126
9], requires_grad=True))
# ...
# ('final.weight', Parameter containing:
# tensor([[-0.0570, 0.4325, 0.4118, -1.6617]], requires_grad=True))
# ('final.bias', Parameter containing:
# tensor([-0.6704], requires_grad=True))到目前为止,我们只定义了模块并调用了它们的 forward() 方法来生成计算输出。为了训练模块,需要使用样本数据并根据该数据调整参数以优化目标函数的结果。
对于以下示例,将使用 PyTorch 的数据加载器接口从随机生成的 DataLoader 中加载样本数据。
from torch.utils.data import DataLoader, TensorDataset
num_samples = 2000
num_features = 10
x = torch.randn(num_samples, num_features)
y = torch.randn(num_samples, 1)
dataset = TensorDataset(x, y)
dataloader = DataLoader(dataset, batch_size=64, shuffle=True)接下来,定义一个简单的神经网络模块。
class SimpleNet(nn.Module):
def __init__(self):
super(SimpleNet, self).__init__()
self.linear1 = nn.Linear(num_features, 5)
self.relu = nn.ReLU()
self.linear2 = nn.Linear(5, 1)
def forward(self, x):
x = self.linear1(x)
x = self.relu(x)
x = self.linear2(x)
return xmodel = SimpleNet()
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)for epoch in range(20):
for batch_x, batch_y in dataloader:
optimizer.zero_grad()
output = model(batch_x)
loss = criterion(output, batch_y)
loss.backward()
optimizer.step()
print(f'Epoch [{epoch+1}/20], Loss: {loss.item():.4f}')这个示例展示了如何构建一个简单的模块并使用 PyTorch 的优化器和损失函数进行训练。通过训练,模型参数将逐渐调整以最小化损失函数的值,从而实现模型对样本数据的最佳拟合。
这样,通过模块和 PyTorch 的各种工具,可以构建、训练和优化复杂的神经网络模型,进而实现各种深度学习任务。
torch.seed()在所有设备上设置用于生成随机数的种子为一个非确定性的随机数。
torch.manual_seed(seed)在所有设备上设置用于生成随机数的种子。
torch.initial_seed()返回用于生成随机数的初始种子。
torch.get_rng_state()返回随机数生成器的状态,类型为torch.ByteTensor。
torch.set_rng_state(state)设置随机数生成器的状态。
torch.default_generator返回默认的CPU torch.Generator。
torch.rand(size)返回一个张量,其中包含从区间 [0, 1) 的均匀分布中生成的随机数。
torch.randint(low, high, size)返回一个张量,其中包含在 [low, high) 区间内均匀生成的随机整数。
torch.randn(size)返回一个张量,其中包含从均值为0、方差为1的正态分布(标准正态分布)中生成的随机数。
torch.randperm(n)返回从0到n-1的随机排列。
torch.bernoulli(input)从伯努利分布中抽取二元随机数(0或1),概率由输入张量的值指定。
torch.multinomial(input, num_samples)返回一个张量,其中每行包含从多项分布(严格定义为多变量分布)中采样的num_samples个索引。
torch.normal(mean, std)返回一个张量,其中包含从均值和标准差指定的正态分布中生成的随机数。
torch.poisson(input)返回一个与输入张量大小相同的张量,其中每个元素是从泊松分布中采样的,速率参数由对应的输入元素指定。
import torch
torch.manual_seed(42)# 生成一个3x3的均匀分布随机张量
rand_tensor = torch.rand(3, 3)
print(rand_tensor)
# 生成一个3x3的标准正态分布随机张量
randn_tensor = torch.randn(3, 3)
print(randn_tensor)
# 生成从0到9的随机排列
randperm_tensor = torch.randperm(10)
print(randperm_tensor)以上列举了在PyTorch中常用的随机采样方法和设置随机数生成器种子的方法。通过合理使用这些方法,可以确保模型训练的可重复性和随机过程的控制。
在PyTorch中,可以使用一些上下文管理器(context managers)来局部禁用或启用梯度计算。这些管理器包括torch.no_grad()、torch.enable_grad()和torch.set_grad_enabled()。这些管理器在本地线程中起作用,因此如果使用threading模块将工作发送到另一个线程,它们将不起作用。
禁用梯度计算的上下文管理器。
启用梯度计算的上下文管理器。
设置梯度计算状态的上下文管理器。
设置梯度计算状态的上下文管理器。
返回当前是否启用了梯度计算。
启用或禁用推理模式的上下文管理器。
返回当前是否启用了推理模式。
import torch
x = torch.zeros(1, requires_grad=True)
with torch.no_grad():
y = x * 2
print(y.requires_grad) # 输出: Falseimport torch
x = torch.zeros(1, requires_grad=True)
# 禁用梯度计算
is_train = False
with torch.set_grad_enabled(is_train):
y = x * 2
print(y.requires_grad) # 输出: False
# 启用梯度计算
torch.set_grad_enabled(True) # 也可以作为一个函数来使用
y = x * 2
print(y.requires_grad) # 输出: True
# 再次禁用梯度计算
torch.set_grad_enabled(False)
y = x * 2
print(y.requires_grad) # 输出: Falseimport torch
torch.set_grad_enabled(True)
print(torch.is_grad_enabled()) # 输出: True
torch.set_grad_enabled(False)
print(torch.is_grad_enabled()) # 输出: Falseimport torch
with torch.autograd.grad_mode.inference_mode():
x = torch.randn(3, 3)
y = x * 2
print(torch.is_inference_mode_enabled()) # 输出: False (因为推理模式只在上下文管理器内有效)以上示例展示了如何使用这些上下文管理器来控制PyTorch中的梯度计算。这些工具对于在训练和推理过程中优化计算资源非常有用。
PyTorch 的分布式数据并行(DDP)模块旨在通过多个 GPU 或机器进行分布式训练。其核心思想是将模型的计算分布在多个设备上,以加快训练过程。关键步骤包括:
- 设置和清理: 使用
setup和cleanup函数初始化和销毁进程组,以实现不同进程之间的通信。 - 模型分布: 使用
DistributedDataParallel(DDP) 将模型复制到每个 GPU 上,确保梯度更新同步。 - 训练循环: 修改训练循环以适应分布式环境,确保每个进程处理部分数据并同步更新。
训练脚本使用 torchrun 命令执行,将工作负载分配到指定数量的 GPU 或节点。
torchrun --nproc_per_node=2 --nnodes=1 example_script.pyAccelerate 是一个轻量级库,旨在简化 PyTorch 代码的并行化过程。它允许在单 GPU 和多 GPU/TPU 设置之间无缝过渡,代码改动最小。主要特点包括:
- Accelerator 类: 处理分布式环境的设置和管理。
- 数据管道效率: 自定义采样器用于优化多个设备的数据加载,减少内存开销。
- 代码简化: 通过
accelerator.prepare封装 PyTorch 组件,使相同代码在任何分布式设置下运行,无需大量修改。
这种方法确保您的训练脚本保持简洁,并在受益于分布式训练能力的同时,保持 PyTorch 的原生结构。
from accelerate import Accelerator
accelerator = Accelerator()
# 使用 accelerator.prepare 准备您的数据加载器、模型和优化器
train_loader, test_loader, model, optimizer = accelerator.prepare(
train_loader, test_loader, model, optimizer
)Hugging Face Trainer API 提供了一个高级接口,用于训练模型,支持各种训练配置,包括分布式设置。它抽象了大量模板代码,让您专注于训练逻辑。主要组件包括:
- TrainingArguments: 配置常见的超参数和训练选项。
- Trainer 类: 处理训练循环、评估和数据加载。您可以子类化 Trainer 以自定义损失计算和其他训练细节。
- 数据整理器: 用于将数据预处理为训练所需的格式。
Trainer API 支持无缝分布式训练,无需大量代码修改,非常适合复杂的训练场景。
from transformers import Trainer, TrainingArguments
training_args = TrainingArguments(
"basic-trainer",
per_device_train_batch_size=64,
per_device_eval_batch_size=64,
num_train_epochs=1,
evaluation_strategy="epoch",
remove_unused_columns=False
)
class MyTrainer(Trainer):
def compute_loss(self, model, inputs, return_outputs=False):
outputs = model(inputs["x"])
target = inputs["labels"]
loss = F.nll_loss(outputs, target)
return (loss, outputs) if return_outputs else loss
trainer = MyTrainer(
model,
training_args,
train_dataset=train_dset,
eval_dataset=test_dset,
data_collator=collate_fn,
)
trainer.train()使用 notebook_launcher,您可以在 Jupyter Notebook 中使用多个 GPU 运行训练脚本。
from accelerate import notebook_launcher
notebook_launcher(train_trainer_ddp, args=(), num_processes=2)DeepSpeed 是由微软开发的一个深度学习优化库,旨在加速和优化大规模模型的训练过程,尤其是在分布式环境中。DeepSpeed 提供了一系列工具和技术,使得训练超大规模的模型(如GPT-3、BERT等)变得更加高效和可扩展。
-
ZeRO(Zero Redundancy Optimizer)优化器: ZeRO 是 DeepSpeed 的核心技术之一,旨在通过减少冗余数据来显著降低显存占用,从而支持超大模型的训练。ZeRO 优化器分为三个阶段:
- ZeRO-1:分布式优化器状态,将优化器状态(如权重、梯度)在多个GPU之间分配,减少每个GPU的内存负担。
- ZeRO-2:分布式梯度计算,将梯度计算的中间结果在多个GPU之间分配,进一步节省显存。
- ZeRO-3:分布式模型参数,将模型参数也在多个GPU之间分配,最大限度地降低每个GPU的内存占用。
-
混合精度训练: DeepSpeed 支持 FP16 混合精度训练,这种方法通过在训练中使用更低的浮点精度(FP16)来加速计算,同时保留了 FP32 的精度进行关键计算。这样可以在不损失模型精度的前提下,显著提升训练速度并降低显存使用。
-
深度模型并行: 除了传统的数据并行和模型并行,DeepSpeed 还支持管道并行,这种方法将模型分为多个阶段,并在不同的设备上并行处理。这种方式特别适用于训练非常深的神经网络。
-
大规模分布式训练: DeepSpeed 通过高效的通信优化和内存管理,使得用户可以在数百甚至数千个 GPU 上进行大规模分布式训练。它集成了诸如 NCCL、Megatron-LM 和 Turing-NLG 等技术,实现了跨 GPU 的高效通信。
-
自动并行化和调优: DeepSpeed 提供了自动化的并行化和超参数调优工具,可以帮助用户轻松设置和优化训练过程。这极大简化了大规模模型训练的难度,使得开发者能够更专注于模型的设计和创新。
- 超大规模语言模型:如 GPT、BERT 等的训练。
- 计算资源受限的环境:在有限的 GPU 资源下训练大模型。
- 快速迭代和实验:通过加速训练过程,提升模型开发效率。
DeepSpeed 是一个强大而灵活的工具,尤其适合需要训练大规模深度学习模型的研究者和工程师。通过其多样化的优化手段,DeepSpeed 可以大幅度降低训练成本,提升训练效率。
在使用 PyTorch 构建神经网络时,理解如何访问和遍历模型的不同组成部分至关重要。PyTorch 提供了一些函数,允许你探索模型中的模块、参数、缓冲区等。包括 modules()、named_buffers()、named_children()、named_modules()、named_parameters() 和 parameters()。
modules() 函数返回一个遍历神经网络中所有模块的迭代器。这包括模型本身及其包含的任何子模块。值得注意的是,重复的模块只会返回一次。
import torch.nn as nn
l = nn.Linear(2, 2)
net = nn.Sequential(l, l)
for idx, m in enumerate(net.modules()):
print(idx, '->', m)输出:
0 -> Sequential(
(0): Linear(in_features=2, out_features=2, bias=True)
(1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)
在这个例子中,即使 Sequential 模块包含两次相同的 Linear 层,modules() 函数在遍历时只返回一次。
named_buffers() 函数返回一个遍历模块中缓冲区的迭代器,返回缓冲区的名称和缓冲区本身。缓冲区是 PyTorch 中的张量,不被视为模型参数(例如,批量归一化层中的运行均值和方差)。
prefix(str):要添加到所有缓冲区名称前的前缀。recurse(bool):如果为True,则包含所有子模块的缓冲区。默认为True。remove_duplicate(bool):是否在结果中移除重复的缓冲区。默认为True。
for name, buf in net.named_buffers():
if name in ['running_var']:
print(buf.size())这个示例展示了如何遍历模型中的缓冲区并根据名称进行过滤。
named_children() 函数提供一个遍历模型直接子模块的迭代器,返回模块的名称和模块本身。
for name, module in net.named_children():
print(name, '->', module)这个函数特别适用于在不深入子模块的情况下,检查或修改模型的特定层。
named_modules() 返回一个遍历网络中所有模块的迭代器,包括子模块,并返回模块的名称和模块本身。与 modules() 类似,此函数只会返回每个模块一次,即使它在网络中出现多次。
memo(Optional[Set[Module]]): 用于存储已添加到结果中的模块的集合。prefix(str): 添加到模块名称前的前缀。remove_duplicate(bool): 是否移除重复的模块实例。
for idx, m in enumerate(net.named_modules()):
print(idx, '->', m)输出:
0 -> ('', Sequential(
(0): Linear(in_features=2, out_features=2, bias=True)
(1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))
在这个例子中,named_modules() 返回的元组包含模块名称和模块本身。
named_parameters() 函数返回一个遍历模块中所有参数的迭代器,返回参数的名称和参数本身。
prefix(str): 要添加到所有参数名称前的前缀。recurse(bool): 如果为True,则包含所有子模块的参数。默认为True。remove_duplicate(bool): 是否移除重复的参数。默认为True。
for name, param in net.named_parameters():
if name in ['bias']:
print(param.size())这个示例遍历模型中的所有参数,并根据名称进行过滤。
最后,parameters() 函数返回一个遍历模块参数的迭代器。这在将参数传递给优化器时尤其有用。
recurse(bool):如果为True,则包含所有子模块的参数。
for param in net.parameters():
print(type(param), param.size())输出:
<class 'torch.Tensor'> torch.Size([2, 2])
<class 'torch.Tensor'> torch.Size([2])
这个函数非常直观,通常在设置模型的优化器时使用。
理解这些函数可以极大地增强你处理复杂 PyTorch 模型的能力。它们提供了灵活的方法来访问和操作网络的不同部分,从模块到参数。通过利用这些函数,你可以更轻松地调试、修改和优化模型。
- 数据加载:DataLoader可以从不同来源加载数据,如硬盘上的文件、数据库、网络等。它能够自动将数据集划分为小批次,从而减小内存需求,确保数据的高效加载。
- 数据批次处理:每个批次由多个样本组成,可以并行地进行数据预处理和数据增强。这有助于提高模型训练的效率,同时确保每个批次的数据都经过适当的处理。
-
根据dataset和sampler,生成数据索引。
-
根据这些索引,从dataset中读取指定数量的数据,并对其进行预处理(例如归一化、裁剪 等)。
-
如果设置了collate_fn,则将处理后的数据打包成批次数据。
-
如果设置了num_workers > 0,则将数据加载任务分配给多个子进程并行完成。
-
在模型训练时,每个epoch从DataLoader中获取一个批次的数据,作为模型的输入。
深度学习框架用计算图来描述模型的拓扑结构。计算图中用节点表示算子,节点间的边表示张量状态。计算图是一个有向无环图,描述算子之间的依赖关系,计算图中要避免循环依赖导致计算锁死,对于循环结构一般进行展开(unrolling)。
计算图可以根据生成方式的不同,分为:静态图和动态图。
静态图是对完整的模型进行编译得到的固定代码文本。可以对静态图进行优化(算子融合等),得到更高效的结构提升硬件计算性能。编译时以数据占位符作为虚拟输入,不进行条件判断,将所有分支算子加入计算图。优势:计算性能,直接部署。
动态图是在执行时(有输入数据)进行利用框架的算子分发功能输出结果,不生成完整的计算图,只有临时的图拓扑结构。在执行的过程中记录算子,张量和梯度信息,前向传播完毕后,串联起来进行反向传播。优势:方便调试,编程友好。
特性对比
| 特性 | 动态图 | 静态图 |
|---|---|---|
| 代码调试 | 灵活,方便 | 固定,不易调试 |
| 模型部署 | 灵活,方便 | 固定,直接部署 |
| 计算性能 | 较低 | 较高 |
| 优化难度 | 较低 | 较高 |
Module是pytorch的基本单元,包括:1、一个构造函数,它为调用准备模块。2、一组参数和子模块。由构造函数初始化,并且可以在调用期间由模块使用。3、正向函数。调用模块时运行的代码。
许多框架采用计算符号导数的方法,给出了完整的模型表示。然而,在PyTorch中使用gradient tape,记录发生的算子,并在计算导数时反向操作。这样,框架就不必为语言中的所有构造明确定义导数。
TorchScript可以从pytorch代码中生成序列化,优化的模型。在python环境下训练好模型,通过torchscrip导出模型,部署到没有python依赖的环境。
转化静态图的意义:1、TorchScript代码可以在它自己的解释器(受限制的python解释器)中调用。此解释器不需要全局解释器锁,因此可以在同一实例上同时处理许多请求。2、这种格式允许我们将整个模型保存到磁盘,并将其加载到另一个环境中,例如在用Python以外的语言编写的服务器中。3、TorchScript为我们提供了一种IR表示,我们可以在其中对代码进行编译器优化,以提供更高效的执行。4、TorchScript允许后端/设备上推理时获取比单个算子更广泛的视图(全局的静态图)
torch.compile就是PyTorch 2.2版本中的一个重要新特性,是一种新的PyTorch编译器,它可以将Python和TorchScript模型编译成TorchDynamo图,从而提高模型的运行效率。
import torch
import torch.compile
# 假设我们有一个简单的PyTorch模型
model = torch.nn.Sequential(
torch.nn.Linear(10, 5),
torch.nn.ReLU(),
torch.nn.Linear(5, 1),
)
# 使用torch.compile将模型编译成TorchDynamo图
compiled_model = torch.compile(model, torch.jit.ScriptModule)
# 现在,我们可以像使用普通的PyTorch模型一样使用compiled_model
input_data = torch.randn(1, 10)
output_data = compiled_model(input_data)
torch.compile相对于之前的PyTorch编译器解决方案,如TorchScript和FX Tracing,有以下几个优势:
- 更灵活的模型定义:与TorchScript相比,torch.compile允许你使用Python直接定义模型,而不需要将模型转换为TorchScript的静态图。这意味着你可以更灵活地定义模型,而不需要考虑TorchScript的限制。
- 更好的性能:TorchDynamo图是一种优化的中间表示形式,它允许PyTorch编译器进行更多的优化,从而提高模型的运行效率。与FX Tracing相比,TorchDynamo图可以提供更好的性能。
- 更易于调试:由于torch.compile允许你使用Python直接定义模型,因此你可以更容易地调试模型。你可以使用Python的调试工具来检查模型的输入和输出,从而更容易地找到和修复错误。
在AI模型的训练过程中,使用可视化工具能够帮助我们直观地观察模型的训练效果、调试超参数以及优化模型。以下是一些经典且实用的训练过程可视化工具:
- 概述:TensorBoard 是 TensorFlow 官方提供的可视化工具,也是目前深度学习中最常用的训练过程可视化工具之一。
- 功能:支持可视化训练指标(如损失、准确率)、查看计算图、参数分布、梯度变化和激活值等。同时也可以进行超参数调优。
- 兼容性:支持 TensorFlow、PyTorch(通过
torch.utils.tensorboard)以及部分其他框架。 - 适用场景:适用于AIGC、传统深度学习、自动驾驶领域AI模型的训练过程分析,尤其适合复杂模型调试。
- 概述:Weights & Biases 是一个功能强大的实验管理和可视化工具,广泛应用于科研和工业界。
- 功能:支持实时监控训练过程、记录和可视化指标、超参数调优、数据版本控制以及生成详细报告。此外,WandB 还可以生成详细的训练报告和可视化模型的权重、梯度等。
- 兼容性:支持多种框架,包括 TensorFlow、Keras、PyTorch、Scikit-Learn 等。
- 适用场景:适合AIGC、传统深度学习、自动驾驶领域中需要跟踪多个实验、超参数搜索和团队协作的项目。
- 概述:MLflow 是一个开源的机器学习实验管理工具,具有模型训练跟踪、项目管理和模型部署等功能。
- 功能:记录训练过程中的指标和参数变化,支持模型版本控制以及跨设备的协作。可以轻松地记录模型运行的结果、训练曲线和模型权重。
- 兼容性:支持 PyTorch、TensorFlow、Scikit-Learn 等。
- 适用场景:适用于AIGC、传统深度学习、自动驾驶领域中希望将可视化和管理结合的场景,尤其是需要进行模型版本控制和跟踪的项目。
- 概述:ClearML 是一个开源的端到端机器学习和深度学习实验管理平台。
- 功能:包括训练监控、任务管理、数据集版本管理和模型部署,提供实时指标可视化、训练曲线、超参数调优支持。
- 兼容性:支持 TensorFlow、PyTorch、Keras 等多种主流深度学习框架。
- 适用场景:ClearML 非常适合AIGC、传统深度学习、自动驾驶领域中需要完整的实验跟踪、管理和可视化解决方案的用户。
- 概述:VisualDL 是百度飞桨(PaddlePaddle)提供的可视化工具,功能与 TensorBoard 类似。
- 功能:支持监控训练指标、展示计算图、参数分布、PR 曲线等。还支持高维数据可视化和超参数调优。
- 兼容性:虽然与 PaddlePaddle 完全兼容,也支持 PyTorch 和 TensorFlow。
- 适用场景:适用于AIGC、传统深度学习、自动驾驶领域中国内开发者,特别是使用飞桨框架的用户。
- 概述:Comet 是一个实验管理和可视化工具,提供了多种模型和实验跟踪功能。
- 功能:支持实时查看训练指标、超参数调优、生成训练曲线、版本管理和协作。还提供模型的可视化及对比功能。
- 兼容性:支持 TensorFlow、Keras、PyTorch、Scikit-Learn 等。
- 适用场景:适合AIGC、传统深度学习、自动驾驶领域中需要强大可视化功能和团队协作的实验项目。
- 概述:Neptune 是一个面向机器学习和深度学习的可视化跟踪平台。
- 功能:支持实时训练过程监控、记录模型参数和指标、超参数调优以及结果共享。Neptune 的最大特色是其与 Jupyter Notebook 深度集成,便于快速调试。
- 兼容性:支持 TensorFlow、PyTorch、Keras 等多个框架。
- 适用场景:适合AIGC、传统深度学习、自动驾驶领域中需要精细化实验管理、跨团队协作和快速调试的用户。
- 概述:Plotly 和 Dash 是用于数据科学和深度学习的强大可视化库,可以实现自定义的实时数据可视化界面。
- 功能:支持创建交互式训练过程图表和统计图,可以结合模型的实时状态做出复杂可视化界面。
- 兼容性:与多种框架兼容,但需要手动集成。
- 适用场景:适合AIGC、传统深度学习、自动驾驶领域中需要高度自定义的可视化、交互式展示的场景。
- 概述:Netron 是一个开源的神经网络模型可视化工具,支持查看模型结构。
- 功能:支持多种深度学习框架和格式的模型可视化,如 TensorFlow、Keras、PyTorch、ONNX 等,主要用于查看模型的各层结构。
- 兼容性:支持多种模型文件格式,包括 ONNX、H5、PB、TFLite 等。
- 适用场景:适合AIGC、传统深度学习、自动驾驶领域中查看模型结构、理解模型架构和调试模型。
从 PyTorch 应用中抓取计算图,相比于 TorchScript 和 TorchFX,TorchDynamo 更加灵活、可靠性更高。TorchScript通过 jit.trace 或者 jit.script 把模型转化为 TorchScript 的过程困难重重,往往需要修改大量源代码。而 TorchFX 在捕获计算图时,遇到不支持的算子会直接报错,最常见的就是 if 语句。TorchDynamo 克服了 TorchScript 和 TorchFX 的缺点,使用起来极为方便,用户体验相比于 TorchScript 和 TorchFX 大幅提升。配合 TorchInductor 等后端编译器,经 TorchDynamo 捕获的计算图只需要几行代码的改动就可以观测到不错的性能提升。
TorchDynamo 捕获计算图是在翻译 Python 字节码的过程中实现的。Python 函数在执行前会被 Python 虚拟机编译为字节码 (bytecode),每一个 Python 函数的实例都对应一个 frame,其中保存着运行该函数所需要的全局变量、局部变量、字节码等等。
TorchDynamo 的 编译过程发生在将要执行前,它是一个 JIT 编译器。在 Python 将要执行函数时,TorchDynamo 开始翻译字节码并捕获计算图。在 Python 虚拟机 (PVM) 中有一个非常重要的函数 _PyEval_EvalFrameDefault,它的功能是在 PVM 中逐条执行编译好的字节码。TorchDynamo 的入口是 PEP-523 提供的 CPython Frame Evaluation API,它可以让用户通过 回调函数(callback function) 获取字节码,并把修改过后的字节码返回给解释器执行,或者执行预先编译好的目标代码,从而可以在 Python 中实现 即时编译器 (JIT Compiler) 的功能。TorchDynamo 正是通过 PEP-523 把 TorchDynamo 的核心逻辑引入到 Python 虚拟机中,从而在函数将要运行前获取字节码。 TorchDynamo 实现了一个 Python 虚拟机的模拟器,在模拟 Python 字节码执行的过程中构建出对应的计算图
- TorchDynamo 的作用是从 PyTorch 程序中捕获计算图;
- TorchDynamo 是一个 JIT compiler,它的工作原理是通过 PEP-523 获取将要执行的 Python 函数的字节码,在翻译字节码的过程中构建 FX Graph;
- 每个编译过的 frame 都有一个 cache,为同一个函数编译的不同输入属性的函数都保存在 cache 中;
- Guard 用来判断是否能够重用已经编译好的函数,它负责检查输入数据的属性有没有发生变化;
- 碰到不支持的算子时,TorchDynamo 会通过 graph break 把计算图切分为子图,不支持的算子由 Python 解释器执行;
- 循环在 TorchDynamo 捕获计算图时被展开;
- TorchDynamo 会试着内联被调函数,如果成功则生成一张大的计算图,失败则在主调函数中创建 graph break;
- TorchDynamo 会在 DDP bucket 的边界引入 graph break,从而确保 allreduce 能与反向传播同时执行;
在 PyTorch 2.0 以前,用户通过 PyTorch 可以直接捕获到正向传播的计算图,比如 JIT trace 和 TorchFX 的 symbolic trace。虽然 PyTorch 的每个算子都包含正向传播和反向传播的实现,但用户并不能直接在反向传播的计算图上面做优化,也无法把正向传播和反向传播的计算图合并在一张计算图中。PyTorch 2.0 中引入了 AOTAutograd,它的出现解决了这个问题,从而使得一些针对 training 的优化变得可能。
有了 AOTAutograd,用户可以做以下事情:
- 获取反向传播计算图、甚至是正向传播和反向传播联合的计算图;
- 用不同的后端编译器分别编译正向传播和反向传播计算图;
- 针对训练 (training) 做正向传播、反向传播联合优化,比如通过在反向传播中重算 (recompute) 来减少正向传播为反向传播保留的 tensor,从而削减内存需求;
PyTorch 反向传播的计算图是在执行正向传播的过程中动态构建的,反向传播的计算图在正向传播结束时才能确定下来。AOTAutograd 以 Ahead-of-Time 的方式同时 trace 正向传播和反向传播,从而在函数真正执行之前拿到正向传播和反向传播的计算图。 工作流程:
- 以 AOT 方式通过 torch_dispatch 机制 trace 正向传播和反向传播,生成联合计算图 (joint forward and backward graph),它是包含 Aten/Prim 算子的 FX Graph;
- 用 partition_fn 把 joint graph 划分为正向传播计算图和反向传播计算图;
- 可选: 通过 decompositions 把 high-level 算子分解、下沉到粒度更小的算子;
- 调用 fw_compiler 和 bw_compiler 分别编译正向传播计算图和反向传播计算图,通过 TorchFX 生成编译后的 Python 代码,并整合为一个 torch.autograd.Function;
- AOTAutograd 利用了 torch_dispatch 机制通过 tracing 提前得到联合正向传播和反向传播计算图;
- 经过 torch_dispatch trace 得到的是最内层的 ATen 算子,AOTAutograd 将其保存在 FX Graph 中;
- 如果用于 tracing 的 tensors 中有重复,那么通过 make_fx 得到的计算图与预期不符,AOTAutograd 会在 tracing 前去重;
- AOTAutograd 用 partition_fn 把 trace 得到的 joint graph 划分为 foward graph 和 backward graph;
- min_cut_rematerialization_partition 通过求解最大流/最小割问题最小化正向传播保留给反向传播的 tensor;
- make_fx 的 tracing 不支持 data-dependent control flow,循环、函数调用在 tracing 后被展开;
在 PyTorch 中,.detach()、.clone()、requires_grad=True 和 torch.no_grad() 是涉及 自动微分(autograd) 和 张量操作 的核心概念。它们控制了张量是否参与计算图的构建、是否跟踪梯度,以及如何高效地操作张量。
.detach() 方法用于从当前的计算图中分离张量。分离后的张量与原张量共享相同的存储空间(数据),但不会再参与梯度计算。
在 PyTorch 的自动微分机制中,每个操作都会在后台构建一个计算图,用于反向传播计算梯度。而 .detach() 会创建一个新的张量,分离计算图:
- 新的张量不跟踪梯度。
- 新张量的
requires_grad属性为False。
- 避免梯度计算:
- 在反向传播中,有些中间结果不需要计算梯度时,使用
.detach()避免冗余计算图构建。
- 在反向传播中,有些中间结果不需要计算梯度时,使用
- 进行无梯度的张量操作:
- 处理某些张量,只需要其值而不需要其与梯度计算的关系。
import torch
# 创建张量并参与计算图
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = x * 2
# 分离张量,y_detached 不再参与计算图
y_detached = y.detach()
# 检查 requires_grad 属性
print(y.requires_grad) # True
print(y_detached.requires_grad) # False
# 修改 y_detached 的值,不影响 y
y_detached[0] = 10
print(y) # tensor([2., 4., 6.], grad_fn=<MulBackward0>).clone() 方法用于深复制一个张量,新张量的存储空间与原张量完全独立。
.clone()创建一个新的张量,具有与原张量相同的数据值和属性(如requires_grad)。- 如果原张量需要梯度,
clone()出的张量会继续参与梯度计算,且其计算图关系保持不变。
- 创建独立的张量拷贝:
- 对原张量的修改不会影响到克隆后的张量。
- 操作过程中需要保留中间结果:
- 尤其在构建复杂的计算图时,使用
.clone()可以保留独立状态。
- 尤其在构建复杂的计算图时,使用
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
# 克隆张量
y = x.clone()
y[0] = 10 # 修改 y 不会影响 x
# 验证
print(x) # tensor([1., 2., 3.], requires_grad=True)
print(y) # tensor([10., 2., 3.], requires_grad=True)
# 克隆张量保持计算图
z = x * 2
z_clone = z.clone()
print(z.grad_fn) # <MulBackward0 object>
print(z_clone.grad_fn) # <MulBackward0 object>张量的 requires_grad 属性控制其是否需要计算梯度。如果设置为 True,该张量会参与计算图的构建,并在反向传播时计算和存储梯度。
- 当
requires_grad=True:- 张量会参与计算图,记录每一步的操作。
- 在反向传播时,PyTorch 会根据计算图,计算梯度并存储在
tensor.grad属性中。
- 当
requires_grad=False:- 张量不会记录操作,且节省内存和计算资源。
- 训练模型时需要梯度:
- 对模型参数(如权重)设置
requires_grad=True,以便在反向传播中更新权重。
- 对模型参数(如权重)设置
- 冻结梯度:
- 在推理阶段,或冻结某些层时,将
requires_grad=False。
- 在推理阶段,或冻结某些层时,将
# 创建需要梯度的张量
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
# 操作
y = x * 2
z = y.sum()
# 反向传播
z.backward()
# 梯度
print(x.grad) # tensor([2., 2., 2.])
# 冻结梯度
x.requires_grad_(False)
print(x.requires_grad) # Falsetorch.no_grad() 是一个上下文管理器,临时禁用自动梯度计算。
在 torch.no_grad() 块内:
- 所有操作都会被标记为不需要梯度。
- 不会构建计算图。
- 可以节省内存和计算资源。
注意:torch.no_grad() 是临时的,只在上下文块中生效。
- 推理阶段:
- 在模型推理中,通常只需要前向传播,使用
torch.no_grad()避免不必要的计算图构建。
- 在模型推理中,通常只需要前向传播,使用
- 冻结梯度操作:
- 修改模型权重或对张量进行操作时,不希望干扰现有计算图。
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
# 不使用 torch.no_grad
with torch.no_grad():
y = x * 2
print(y.requires_grad) # False
# 离开 no_grad 块后
z = x * 2
print(z.requires_grad) # True| 功能 | 作用 | 是否构建计算图 | 是否跟踪梯度 |
|---|---|---|---|
.detach() |
分离张量与计算图,不再跟踪梯度 | 否 | 否 |
.clone() |
深复制张量,生成新张量(可继续跟踪梯度) | 是(如需要) | 是(如需要) |
requires_grad=True |
控制张量是否需要计算梯度 | 是 | 是 |
torch.no_grad() |
上下文管理器,禁用梯度计算(节省资源) | 否 | 否 |
- 训练阶段:
requires_grad=True保证计算图的正确构建。 - 推理阶段:使用
torch.no_grad(),避免不必要的计算图构建。 - 冻结部分梯度:使用
.detach()或设置requires_grad=False。 - 深复制张量:使用
.clone()确保独立性。
在 PyTorch 中,连续张量(contiguous tensor)和非连续张量(non-contiguous tensor)的区别主要涉及内存布局和访问方式。了解这些区别对于我们在AIGC、传统深度学习以及自动驾驶中高效地操作张量和调试潜在的错误非常重要。
- 连续张量:内存布局是线性的,操作高效且兼容性好。
- 非连续张量:内存布局非线性,可能导致性能开销和潜在错误。
- 可以使用
is_contiguous()检查张量连续性,用contiguous()将非连续张量转换为连续张量。
- 连续张量的内存布局是线性的,即数据在内存中按行优先(C 风格)顺序存储,没有跳跃。
- 张量的每个元素的内存地址是相邻的。
- 直接通过
torch.Tensor.contiguous()检查某个张量是否是连续的。
例如,对于一个 3x3 张量:
tensor = torch.tensor([[1, 2, 3],
[4, 5, 6],
[7, 8, 9]])内存布局是连续的:
1 2 3 4 5 6 7 8 9
- 非连续张量的内存布局可能是非线性的,通常由于操作(如
transpose,permute等)导致张量的内存布局变得复杂。 - 元素的存储顺序可能不再按行优先顺序排列。
- 尽管张量的形状看起来相同,但实际内存布局可能不同,访问非连续张量时需要额外的开销。
例如,对张量执行转置操作后:
tensor = torch.tensor([[1, 2, 3],
[4, 5, 6],
[7, 8, 9]])
transposed_tensor = tensor.t() # 转置转置后的内存布局仍是原始的(按行存储的):
1 2 3 4 5 6 7 8 9
但访问顺序变为列优先,这种布局是非连续的。
可以通过 is_contiguous() 方法检查张量是否是连续的:
tensor = torch.tensor([[1, 2, 3],
[4, 5, 6],
[7, 8, 9]])
print(tensor.is_contiguous()) # True
transposed_tensor = tensor.t()
print(transposed_tensor.is_contiguous()) # False- 计算效率高:许多底层操作依赖连续的内存布局,处理连续张量更高效。
- 与第三方库的兼容性:某些操作或库(如 cuDNN)需要输入张量是连续的,否则可能报错。
- 性能开销:在使用非连续张量时,某些操作可能需要额外的内存复制(
contiguous()操作)来重新组织内存布局。 - 潜在错误:某些函数可能要求输入张量是连续的,如果输入非连续张量,可能会报错。
如果需要将非连续张量转换为连续张量,可以使用 contiguous() 方法:
tensor = torch.tensor([[1, 2, 3],
[4, 5, 6],
[7, 8, 9]])
transposed_tensor = tensor.t() # 转置,非连续
print(transposed_tensor.is_contiguous()) # False
contiguous_tensor = transposed_tensor.contiguous()
print(contiguous_tensor.is_contiguous()) # True注意:
contiguous()会创建一个新的张量,并将数据拷贝到新的内存区域,重新排列成连续布局。- 如果张量已经是连续的,
contiguous()不会产生额外开销。
| 操作 | 连续性结果 |
|---|---|
transpose |
非连续 |
permute |
非连续 |
reshape |
通常连续(除非需要与原数据共享内存) |
contiguous |
连续 |
Triton是OpenAI 研发的一个专门为深度学习和高性能计算任务设计的编程语言和编译器,它旨在简化并优化在GPU上执行的复杂操作的开发。Triton 的目标是提供一个开源环境,以比 CUDA 更高的生产力编写快速代码。传统的基于 CUDA 进行 GPU 编程难度较大,学术界和工业界都对面向 GPU 编程的领域特定语言(DSL)很感兴趣。但是目前已有的 DSL 在灵活性和(对相同算法)速度上明显慢于像 cuBLAS、cuDNN 或 TensorRT 这样的库中可用的最佳手写计算内核。已有的 DSL 如 polyhedral machinery (Tiramisu/Tensor Comprehensions)、scheduling languages (Halide、TVM) 等在效率上还有提升空间。 Triton 被提出就是希望作为一个编写灵活的 DSL 来将低 GPU 编程难度的同时提升也提升算子效率
Triton 的核心理念是基于分块的编程范式可以促进神经网络的高性能计算核心的构建。CUDA 编写属于传统的 “单程序,多数据” GPU 执行模型,在线程的细粒度上进行编程,Triton 是在分块的细粒度上进行编程。例如,在矩阵乘法的情况下,CUDA和Triton有以下不同
可以看出 triton 在循环中是逐块进行计算的。这种方法的一个关键优势是,它导致了块结构的迭代空间,相较于现有的DSL,为程序员在实现稀疏操作时提供了更多的灵活性,同时允许编译器为数据局部性和并行性进行积极的优化。
- 易用性:Triton 提供了一个易于使用的编程接口,使得开发者可以轻松地编写高效的 GPU 代码。
- 灵活性:Triton 支持多种编程范式,包括命令式编程和声明式编程,使得开发者可以根据自己的需求选择最合适的编程方式。
- 性能:Triton 的编译器能够自动优化代码,生成高效的 GPU 内核,从而提高计算性能。
import triton
from triton import language as tl
@triton.jit
def add_kernel(
in_ptr0,
in_ptr1,
out_ptr,
n_elements,
BLOCK_SIZE: "tl.constexpr",
):
pid = tl.program_id(axis=0)
block_start = pid * BLOCK_SIZE
offsets = block_start + tl.arange(0, BLOCK_SIZE)
mask = offsets < n_elements
x = tl.load(in_ptr0 + offsets, mask=mask)
y = tl.load(in_ptr1 + offsets, mask=mask)
output = x + y
tl.store(out_ptr + offsets, output, mask=mask)
@torch.compile(fullgraph=True)
def add_fn(x, y):
output = torch.zeros_like(x)
n_elements = output.numel()
grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=4)
return output
x = torch.randn(4, device="cuda")
y = torch.randn(4, device="cuda")
out = add_fn(x, y)
print(f"Vector addition of\nX:\t{x}\nY:\t{y}\nis equal to\n{out}")这段代码是一个用于执行向量加法的 Triton 内核定义,使用 @triton.jit 装饰器进行即时编译(JIT)以便在 GPU 上执行。它逐元素地将两个向量相加,并将结果存储在第三个向量中。
- 输入和输出指针:in_ptr0, in_ptr1, 和 output_ptr 分别是指向第一个输入向量、第二个输入向量和输出向量的指针。这些向量存储在 GPU 的内存中。
- 向量大小和块大小:n_elements 是向量中元素的总数。BLOCK_SIZE 是一个编译时常量(tl.constexpr),定义了每个 GPU 程序(或称为线程块)应该处理的元素数量。
- 程序标识和数据偏移:通过 tl.program_id(axis=0) 获取当前程序(线程块)的唯一标识符 pid。然后,根据这个 ID 和块大小计算出这个程序负责处理的数据段的起始偏移量 block_start。每个程序负责处理一小块数据,这样可以并行处理整个向量。
- 内存访问和掩码:offsets 计算每个元素在向量中的位置,mask 用于创建一个布尔掩码,以防止对数组界外的内存进行操作。这是必要的,因为向量的大小可能不是块大小的整数倍,导致最后一个程序块可能没有足够的元素来处理。
- 加载、计算和存储:使用 tl.load 函数根据 offsets 和 mask 从输入向量中加载元素,执行加法操作得到 output,然后再使用 tl.store 将计算结果根据相同的 offsets 和 mask 存储回输出向量。 这种方法利用了 GPU 的并行计算能力,通过将数据分块并分配给多个程序(线程块)来加速向量加法操作。通过适当选择 BLOCK_SIZE,可以优化内核的性能,以适应特定的硬件和问题规模。
- triton.jit:用于使用Triton编译器对函数进行JIT编译,该函数将在GPU上编译运行,使用jit装饰器的函数只能访问Python基础数据类型、Triton包内的内置函数、该函数的参数以及其他JIT函数。
- triton.autotune:用于评估所有配置,Kernel将会被运行多次后选择最优的配置进行执行,常见的配置属性包括num_warps、num_stages以及BLOCK_SIZE等,具体用法示例如下。
@triton.autotune(configs=[
triton.Config(kwargs={'BLOCK_SIZE': 128}, num_warps=4),
triton.Config(kwargs={'BLOCK_SIZE': 1024}, num_warps=8),
],
key=['x_size'] # the two above configs will be evaluated anytime
# the value of x_size changes
)
@triton.jit
def kernel(x_ptr, x_size, **META):
BLOCK_SIZE = META['BLOCK_SIZE']- triton.heuristics:类似于triton.autotune,不同在于这种方法允许根据输入参数动态计算配置参数,从而提供Triton Kernel的灵活性。
- triton.Config:表示triton.autotune要评估的配置参数。
- Programming Model编程模型:张量tl.tensor、获取给定轴上当前线程块ID tl.program_id……
- Creation Ops:指定[start, end)范围内的连续值tl.arange、填充指定形状和数据类型标量值的张量tl.full、类型转换tl.cast……
- Shape Manipulation Ops形状变换:广播tl.broadcast、维度扩充tl.expand_dims、转置tl.trans、tl.reshape……
- 矩阵乘Ops:tl.dot
- 内存操作:加载tl.load、存储tl.strore……
- 索引操作:tl.where、tl.swizzle2d……
- Math Ops数学库操作:取绝对值tl.abs、向上取整除法tl.cdiv、向上取整tl.ceil、余弦函数tl.cos、正弦函数tl.sin、分类问题中的激活函数tl.softmax、平方根tl.sqrt……
- Reduction Ops归约操作:tl.max、tl.min、tl.reduce、tl.sum……
- Atomic Ops原子操作:tl.atomic_add、tl.atomic_max、tl.atomic_min……
- 随机数生成:tl.randint、tl.rand、tl.randn
- 迭代器:tl.range、tl.static_range
- Debug Ops调试函数:编译时打印tl.static_print、编译时断言tl.static_assert、运行时打印tl.device_print、运行时断言tl.device_print(注意:如果直接在Triton程序中使用Python的打印函数,会在编译运行时直接报错)
triton编译架构
前端(Frontend): 用于将用户使用Python编写的kernel或者Pytorch 2.0中通过Inductor生成的TritonKernel转换为对应的Triton IR
优化器(Optimizer):通过各类pass将Triton IR逐步转换并优化为TritonGPU IR
后端(Backend):将TritonGPU IR逐步转换为LLVM IR,对于NVIDIA GPU最终会被编译为cubin
程序首先被转换成 Triton IR,然后根据目标硬件的特点进行优化,转换为 Triton GPU IR;优化后的 TritonGPU IR 可以转换成 LLVM IR,然后利用 LLVM 的强大优化和代码生成能力,生成可在目标硬件上运行的高效代码。
Triton IR 是 Triton 编译器的高级中间表示,用于表示深度学习模型的计算图,并且是硬件无关的。Triton IR 有以下几个特点:
- 高级抽象:Triton IR 提供了高级抽象,使得开发者可以用接近于高级深度学习框架的方式来描述计算图。
- 操作表示:它包含了一系列的操作(ops),如矩阵乘法、卷积、激活函数等,这些都是深度学习模型中常见的操作。
- 优化:在 Triton IR 层面,编译器可以应用一些高级优化,如死代码消除、常量折叠等。
- 转换:Triton IR 可以被转换为更接近硬件的 Triton GPU IR,以便进行进一步的优化。
Intermediate Representation(IR,中间表示)是在计算机科学和编译器设计中使用的概念,它是一种中间形式的程序表示,用于在不同编译阶段之间传递和处理代码。下面将详细介绍Intermediate Representation的背景、特点和用途。
-
语法分析:编译器将源代码解析为IR形式,以便进行后续的语义分析和优化。
-
优化:IR提供了一种通用的表示形式,使得编译器能够对程序进行各种优化,如常量折叠、死代码消除、循环优化等。
-
中间代码生成:IR可以作为源代码和机器代码之间的桥梁,编译器可以将IR转换为特定平台的机器代码。
-
跨平台编译:通过使用IR,可以实现在不同平台上的代码移植和交叉编译,促进软件的可移植性和跨平台性。
-
混合语言编程:将不同语言的代码转换为共同的IR形式,使得不同语言之间的互操作性成为可能。
在AI编译中,IR是一种中间表示形式,用于表示深度学习模型的计算图。它是一种抽象层次介于源代码和机器代码之间的表示形式,可以用于优化和转换深度学习模型。
作用:
- 跨平台兼容性:IR 提供了一种与硬件无关的表示,使得同一个模型能够适配不同硬件。模型经过编译成 IR 后,可以针对不同硬件进行进一步优化和转换。
- 结构化优化基础:IR 是图优化的基础,因为它提供了模型计算的清晰描述,编译工具可以基于它对模型执行的顺序、内存分配等进行优化。
类型:
- 静态 IR:如 ONNX、TensorFlow GraphDef,这种 IR 表达了完整的模型结构,可用于静态推理和优化。
- 动态 IR:一些框架如 PyTorch 使用动态图 IR,可以在模型推理时灵活调整图结构,适合处理动态输入等场景。
计算图作为连接深度学习框架和前端语言的主要中间表达,被目前主流框架如TensorFlow和PyTorch所使用或者作为标准文件格式来导出模型。 计算图是一个有向无环图(DAG),节点表示算子,边表示张量或者控制边(control flow),节点之间的依赖关系表示每个算子的执行顺序。
计算图的优化被定义为作为在计算图上的函数,通过一系列等价或者近似的优化操作将输入的计算图变换为一个新的计算图。其目标是通过这样的图变换来化简计算图,从而降低计算复杂度或内存开销。在深度神经网络编译器中,有大量优化方法可以被表示为计算图的优化,包括一些在传统程序语言编译器中常用的优化方法。
-
算术表达式化简:,在计算图中的一些子图所对应的算术表达式,在数学上有等价的化简方法来简化表达式,这反应在计算图上就是将子图转化成一个更简单的子图(如更少的节点),从而降低计算量
-
公共子表达式消除(Common Subexpression Elimination, CSE:通过找到程序中等价的计算表达式,然后复用结果,消除其它冗余表达式的计算。同理,在计算图中,公共子表达式消除就等同于寻找并消除冗余的计算子图。一个简单的实现算法是按照图的拓扑序(保证一个访问一个节点时,其前继节点均已经访问)遍历图中节点,每个节点按照输入张量和节点类型组合作为键值进行缓存,后续如果有节点有相同的键值则可以被消除,并且将其输入边连接到缓存的节点的输入节点上。
-
常数传播(constant propagation):也叫常数折叠(constant folding),其主要方法是通过在编译期计算出也是常数表达式的值,用计算出的值来替换原来的表达式,从而节省运行时的开销。在计算图中,如果一个节点的所有输入张量都是常数张量的话,那么这个节点就可以在编译期计算出输入张量,并替换为一个新的常数张量
-
矩阵乘自动融合:矩阵乘在深度学习计算图中被广泛应用,如常见的神经网络的线性层、循环神经网络的单元层、注意力机制层等都有大量的矩阵乘法。在同一个网络里,经常会出现形状相同的矩阵乘法,根据一些矩阵的等价规则,如果把些矩阵乘算子融合成一个大的矩阵乘算子,可以更好的利用到GPU的算力,从而加速模型计算
-
算子融合:在深度学习模型中,针对大量的小算子的融合都可以提高GPU的利用率,减少内核启动开销、减少访存开销等好处。例如,Element-wise的算子(如Add,Mul,Sigmoid,Relu等)其计算量非常小,主要计算瓶颈都在内存的读取和写出上,如果前后的算子能够融合起来,前面算子的计算结果就可以直接被后面算子在寄存器中使用,避免数据在内存的读写,从而提交整体计算效率。
-
子图替换和随机子图替换:鉴于算子融合在深度学习计算中能够带来较好的性能优化,然而在实际的计算图中有太多算子无法做到自动的算子融合,主要原因包括算子的内核实现逻辑不透明、算子之前无法在特有加速器上融合等等。为了在这些的情况下还能进行优化,用户经常会实现一些手工融合的算子来提升性能。那么,编译器在计算图中识别出一个子图并替换成一个等价的新的算子或子图的过程就是子图替换优化。
钩子函数(Hook Function)是AI行业中实现模型可解释性和可控性的核心技术。其核心原理类似于在神经网络的关键位置安装可编程传感器,当数据流经这些节点时自动触发预设操作。这种机制实现了非侵入式监控——无需修改模型结构,即可动态捕捉、干预或记录模型内部状态。
通过合理运用钩子函数,AI行业开发者可获得堪比"神经内窥镜"的模型洞察力,是提升AI系统可靠性和透明度的关键利器。
-
触发机制:
- 前向钩子:在张量通过某层时触发(如
forward_hook) - 反向钩子:在梯度反向传播时触发(如
backward_hook)
- 前向钩子:在张量通过某层时触发(如
-
核心功能:
def hook_function(module, input, output): # 可执行操作:记录、修改、可视化等 modified_output = output * 0.5 # 示例:修改输出 return modified_output
-
生命周期:
- 注册:
layer.register_forward_hook(hook_function) - 执行:随数据流自动触发
- 销毁:通过句柄主动移除
- 注册:
想象一个智能快递分拣系统:
-
原始流程: 包裹 → 扫码区 → 重量检测 → 分拣出口
-
添加钩子:
- 在扫码区后安装摄像头(前向钩子): 记录包裹外观,统计分类错误
- 在重量检测前加装调节器(反向钩子): 对超重包裹自动减重后再检测
类比说明:
- 摄像头:记录中间特征(如ResNet某层激活值)
- 调节器:修改梯度流动(如梯度裁剪)
- 问题:图像生成出现肢体畸变
- 钩子方案:
def decoder_hook(module, input, output): # 在U-Net解码层监测手部区域激活值 hand_mask = create_hand_mask(output.shape) output[hand_mask] *= 1.2 # 增强手部细节 return output unet.decoder[4].register_forward_hook(decoder_hook)
- 效果:寻找与手部生成相关的参数,进行优化增强,手部生成准确率提升37%,无需重新训练模型
- 问题:注意力机制失效导致分类错误
- 钩子方案:
attention_scores = [] def attn_hook(module, input, output): scores = output[1] # 获取注意力权重 attention_scores.append(scores.detach()) transformer.blocks[2].attn.register_forward_hook(attn_hook)
- 分析:可视化第3层注意力头,发现90%的head关注了停用词,从而进行针对性优化训练
- 问题:激光雷达点云误识别导致急刹
- 钩子方案:
def safety_hook(module, input, output): if output[:, 'pedestrian'] > 0.7: send_alert() # 触发安全警报 return output * 0 # 阻断危险信号 return output detection_head.register_forward_hook(safety_hook)
- 效果:假设通过钩子函数实时监测,来减少误检引发的急刹事件。
| 功能 | 传统方法 | 钩子函数方案 |
|---|---|---|
| 中间层访问 | 需修改模型结构 | 非侵入式动态接入 |
| 实时干预 | 几乎不可行 | 毫秒级响应延迟 |
| 计算开销 | 重新训练增加300%耗时 | 额外增加<5%推理时间 |
| 部署灵活性 | 需重新导出模型 | 热插拔式启用/禁用 |
- 性能优化:
- 使用
torch.jit.script编译钩子函数 - 异步处理非关键监控任务
- 使用
- 安全防护:
try: output = module(input) except HookException as e: activate_fallback_model() # 应急系统启动
- 调试工具:
- PyTorch的
torch.utils.hooks模块 - TensorFlow的
tf.keras.callbacks.LambdaCallback
- PyTorch的
PyTorch 的 DataLoader 是数据加载的核心工具,其设计目标是高效、灵活地管理大规模数据集,在 AIGC、传统深度学习、自动驾驶三大领域中,DataLoader 通过不同的配置模式,成为支撑大规模模型训练与实时推理的基石。
-
Dataset
- 定义数据的存储结构和单样本访问接口(
__getitem__和__len__)。 - 示例:图像分类任务中,每个样本对应一个图像文件路径和标签。
- 定义数据的存储结构和单样本访问接口(
-
Sampler
- 控制数据的采样顺序(如顺序采样、随机采样、分布式采样)。
- 生成索引序列(indices),决定 DataLoader 如何遍历数据集。
-
DataLoaderIter
- 内部迭代器,负责从 Sampler 获取索引,从 Dataset 加载数据,并合并为批次(Batch)。
- 支持多进程加速(通过
num_workers参数)。
-
Collate Function
- 将多个单样本数据(如多个图像张量)合并为一个批次张量。
- 默认支持张量拼接,可自定义处理复杂数据结构。
以下通过 图像分类任务 的案例,逐步解析 DataLoader 的运作机制:
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
class ImageDataset(Dataset):
def __init__(self, image_paths, labels, transform=None):
self.image_paths = image_paths # 图像路径列表
self.labels = labels # 标签列表
self.transform = transform # 数据增强(如随机裁剪、归一化)
def __len__(self):
return len(self.image_paths)
def __getitem__(self, idx):
image = Image.open(self.image_paths[idx])
label = self.labels[idx]
if self.transform:
image = self.transform(image)
return image, label
dataset = ImageDataset(paths, labels, transform=...)
dataloader = DataLoader(
dataset,
batch_size=32,
shuffle=True,
num_workers=4,
collate_fn=lambda batch: (torch.stack([x[0] for x in batch]),
torch.tensor([x[1] for x in batch]))-
Sampler 生成索引
shuffle=True时,生成随机索引序列(如[5, 2, 7, ..., 1])。- 每个 epoch 重新生成索引,保证数据随机性。
-
多进程数据加载
- 主进程创建
num_workers个子进程。 - 每个子进程预加载
prefetch_factor * batch_size个样本到内存队列。
- 主进程创建
-
数据合并与返回
- 主进程从队列中取出多个样本,调用
collate_fn合并为批次张量。 - 最终输出格式:
(batch_images, batch_labels)。
- 主进程从队列中取出多个样本,调用
-
多进程加速原理
- 子进程通过共享内存(Shared Memory)传递数据,避免进程间拷贝开销。
- 通过
torch.multiprocessing实现跨进程张量传输。
-
内存管理
- 使用内存池(Memory Pool)缓存常用数据(如归一化后的图像张量)。
- 通过
pin_memory=True启用锁页内存,加速 CPU→GPU 数据传输。
-
动态批处理
- 支持可变长度输入(如文本序列),通过自定义
collate_fn动态填充(Padding)。
- 支持可变长度输入(如文本序列),通过自定义
- 应用场景:训练 Stable Diffusion 生成高分辨率图像。
- DataLoader 优化:
dataloader = DataLoader( dataset, batch_size=4, num_workers=8, # 高并行度加速海量数据加载 pin_memory=True, # 加速数据到 GPU 的传输 persistent_workers=True # 保持子进程存活,避免重复初始化 )
- 技术价值:
单 GPU 训练时,DataLoader 的多进程预加载可隐藏图像解码(如 PNG→Tensor)耗时,使 GPU 利用率保持在 95% 以上。
- 应用场景:ImageNet 大规模分类任务。
- DataLoader 设计:
dataloader = DataLoader( dataset, batch_size=256, sampler=DistributedSampler(dataset), # 分布式训练 collate_fn=custom_collate, # 处理不同尺寸图像 )
- 技术价值:
分布式 Sampler 确保每个 GPU 处理不重叠的数据分区,结合多进程加载,实现线性加速比。
- 应用场景:实时处理激光雷达点云数据。
- DataLoader 扩展:
class LidarDataset(Dataset): def __getitem__(self, idx): point_cloud = load_pcd(self.paths[idx]) # 加载点云文件(.pcd) return voxelize(point_cloud) # 体素化处理 dataloader = DataLoader( LidarDataset(), batch_size=32, num_workers=4, collate_fn=pad_voxels # 动态填充不同数量的体素 )
- 技术价值:
动态批处理支持非均匀点云输入,结合 CUDA 加速的体素化操作,满足实时推理需求(<100ms 延迟)。
特点:
- 灵活性与工业部署:支持静态图与动态图混合编程(Eager Execution),提供SavedModel格式实现跨平台部署(移动端、嵌入式设备)。
- 生态系统完善:集成TensorBoard可视化工具、TFX流水线工具,支持分布式训练和TPU加速。
- 多语言支持:核心为C++,上层API支持Python、Java等,适合大规模生产环境。
实际案例:
- 医疗影像分类:使用预训练的ResNet模型(TensorFlow Hub),在肺部CT图像中快速实现肺癌检测,准确率提升15%。
领域应用:
- AIGC:Stable Diffusion的底层优化中,利用TensorFlow Lite将模型压缩至移动端,实时生成高清图像。
- 传统深度学习:ImageNet图像分类任务中,通过Keras API快速构建InceptionV3模型。
- 自动驾驶:Waymo使用TensorFlow部署多传感器融合模型,实时处理激光雷达与摄像头数据。
特点:
- 动态计算图:支持即时调试,适合研究场景;Autograd机制简化反向传播实现。
- 社区与科研支持:学术界主流框架,提供TorchScript实现模型序列化,兼容ONNX格式。
- GPU加速优化:原生支持CUDA,混合精度训练效率提升30%。
实际案例:
- 自然语言处理:基于Transformer架构(Hugging Face库),训练GPT-3模型实现文本生成,参数量达1750亿。
领域应用:
- AIGC:Meta的Make-A-Video工具使用PyTorch动态图特性,实现文本到视频的跨模态生成。
- 传统深度学习:Fast.ai库简化图像分类任务,5行代码完成CIFAR-10数据集训练。
- 自动驾驶:特斯拉Autopilot采用PyTorch构建BEV(鸟瞰图)感知模型,优化多目标跟踪算法。
特点:
- 高层抽象与易用性:提供模块化API(如
Sequential和Functional),10分钟内可搭建完整神经网络。 - 多后端支持:兼容TensorFlow、Theano等,适合快速原型开发。
- 性能局限:过度封装导致灵活性不足,大规模训练速度较慢。
实际案例:
- 手写数字识别:通过MNIST数据集训练LeNet-5模型,代码量仅20行,准确率达99%。
领域应用:
- AIGC:快速试验生成对抗网络(GAN)结构,如DCGAN生成动漫头像。
- 传统深度学习:教育领域用于教学演示,如LSTM时间序列预测。
- 自动驾驶:小规模车载系统(如停车标志检测)的轻量级模型部署。
特点:
- 高效与模块化:专精计算机视觉,模型Zoo提供预训练权重(如AlexNet、VGG)。
- 工业级优化:支持单GPU每日处理6000万张图像,延迟低至1ms/帧。
- 扩展性差:新增网络层需手动实现C++/CUDA代码,对RNN支持薄弱。
实际案例:
- 图像分类:2014年ImageNet冠军模型GoogLeNet基于Caffe实现。
领域应用:
- AIGC:较少使用,但在早期风格迁移(如Neural Style)中有应用。
- 传统深度学习:工业质检中快速部署ResNet进行缺陷检测。
- 自动驾驶:Mobileye使用Caffe2优化车载视觉模型,实现实时车道线识别。
特点:
- 分布式性能:支持多GPU/多节点训练,内存占用优化(较TensorFlow减少20%)。
- 多语言接口:兼容Python、Scala、Julia,适合云平台(如AWS)。
- 文档不足:社区活跃度低,新手学习成本高。
实际案例:
- 推荐系统:亚马逊使用MXNet构建深度因子分解机(DeepFM),提升广告点击率预测精度。
领域应用:
- AIGC:结合Apache TVM编译器,优化Stable Diffusion推理速度。
- 传统深度学习:大规模NLP任务(如BERT预训练)的分布式加速。
- 自动驾驶:百度Apollo早期版本采用MXNet处理高精地图生成。
- PaddlePaddle:百度开发的国产框架,侧重工业落地,支持模型压缩(如PaddleSlim)。
- JAX:谷歌科研框架,基于自动微分和XLA编译器,适合高性能计算(如AlphaFold)。
- 科研与快速迭代:优先选择PyTorch(动态图+社区资源)。
- 工业部署与跨平台:TensorFlow(生态系统+生产工具链)。
- 轻量级原型开发:Keras(易用性+快速验证)。
- 计算机视觉专项:Caffe/Caffe2(性能优化+预训练模型)。
通过结合具体场景需求与框架特性,我们可最大化AIGC技术方案的部署效率与推理效果。
-
内存访问模式低效
Caffe默认采用NCHW数据布局,Depthwise卷积的逐通道计算导致内存访问跳跃性高,无法充分利用缓存局部性。例如,输入尺寸为[N, C, H, W]时,每个通道独立处理,导致相邻内存地址的数据无法批量加载。 -
缺乏专用算法优化
- 标准卷积通常使用im2col + GEMM优化,但Depthwise的通道间无交互,GEMM效率低下。
- 未应用Winograd算法(可减少乘加运算量),导致计算密度不足。
-
并行化不足
Caffe的OpenMP线程调度在逐通道处理时粒度较粗,无法有效利用多核CPU。例如,当通道数C=32时,可能仅分配4线程,导致负载不均。 -
指令集未充分使用
未针对ARM NEON或Intel AVX2进行汇编级优化,例如未展开循环或使用SIMD向量化指令。
案例验证:
在MobileNetV1的Depthwise卷积层(3x3 kernel,输入尺寸112x112x32)中,Caffe原始实现耗时15ms,而TensorFlow Lite优化后仅需3ms。
-
算法级优化
-
Winograd变换:将3x3卷积转换为4x4矩阵运算,减少60%乘法操作。
$$F(4,3) \text{ Winograd}: 需要16次乘法(原需3x3=9次/点,总16x9=144 → 16x16=256)$$ -
直接卷积优化:避免im2col,通过滑动窗口直接计算,减少内存占用。
-
-
数据布局调整
将NCHW转为NHWC,提升内存连续性。例如,输入[N, H, W, C]更适应Depthwise的逐通道处理模式。 -
指令集加速
- ARM NEON:使用SIMD指令并行处理4个通道(float32x4_t)。
- CPU多线程:细粒度划分任务,如按行或通道块分配线程。
-
集成加速库
调用Intel MKL-DNN或NVIDIA CuDNN的Depthwise专用API(如CuDNN的cudnnDepthwiseConvolutionForward)。
优化效果:
经上述优化后,同一MobileNetV1层的推理时间从15ms降至2.5ms,提升6倍。
-
定义层参数(Proto)
在caffe.proto中新增层的参数消息,例如添加MyLayerParameter:message LayerParameter { optional MyLayerParameter my_layer_param = 12345; } message MyLayerParameter { optional float scale = 1 [default = 1.0]; }
-
实现层类(C++)
- 头文件
my_layer.hpp:继承Layer类,声明Forward/Backward函数。 - 源文件
my_layer.cpp:实现CPU/GPU逻辑,例如:template <typename Dtype> void MyLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom, const vector<Blob<Dtype>*>& top) { const Dtype* input = bottom[0]->cpu_data(); Dtype* output = top[0]->mutable_cpu_data(); caffe_cpu_scale(count, scale_, input, output); }
- 头文件
-
注册层工厂
在layer_factory.cpp中添加注册宏:INSTANTIATE_CLASS(MyLayer); REGISTER_LAYER_CLASS(My);
-
编译与测试
- 修改
Makefile或CMakeLists.txt,包含新层源码。 - 编写测试用例,验证数值正确性(如梯度检查)。
- 修改
案例:添加ScaleLayer实现通道加权,用于风格迁移模型的细节增强。
-
AIGC(生成式AI)
- 应用场景:移动端实时风格迁移
- 技术实现:
优化Depthwise卷积加速轻量级GAN模型(如MobileStyleGAN),在Snapchat等App中实现60FPS滤镜生成。
-
传统深度学习
- 应用场景:工业质检
- 技术实现:
在Caffe中新增DefectDetectLayer,融合Depthwise优化后的MobileNetV2,实现芯片缺陷检测(吞吐量提升3倍)。
-
自动驾驶
- 应用场景:车载实时语义分割
- 技术实现:
使用优化后的Depthwise卷积构建轻量级DeepLabV3+,在NVIDIA Jetson平台达到30FPS,添加BilinearUpsampleLayer提升边缘精度。
torchrun 是 PyTorch 提供的分布式训练启动工具,用于简化多进程或多节点训练任务的配置和管理。它替代了旧版的 torch.distributed.launch,支持弹性训练(自动处理节点故障)和更灵活的参数配置。
--nproc_per_node:每个节点的进程数(通常等于GPU数量)。--nnodes:总节点数(单机为1,多机需指定)。--node_rank:当前节点的序号(主节点为0,其他节点依次递增)。--master_addr:主节点的IP地址。--master_port:主节点的通信端口(默认29500)。
- 初始化进程组:每个进程通过
torch.distributed.init_process_group初始化通信。 - 封装模型:使用
DistributedDataParallel(DDP)包装模型,实现数据并行。 - 数据分片:使用
DistributedSampler对数据集分片,确保各进程处理不同数据。
在单机4卡GPU上训练ResNet-50模型,使用CIFAR-10数据集。
# 单机4卡启动
torchrun --nproc_per_node=4 train.py- 每个进程自动分配独立的GPU(通过
LOCAL_RANK环境变量)。
- 应用场景:训练大规模扩散模型(如Stable Diffusion)。
- 案例:在4机32卡集群中训练文本到图像生成模型。
# 节点0(主节点) torchrun --nnodes=4 --nproc_per_node=8 --node_rank=0 --master_addr=192.168.1.100 train.py # 节点1-3 torchrun --nnodes=4 --nproc_per_node=8 --node_rank=1 --master_addr=192.168.1.100 train.py
- 优势:分布式训练加速迭代,支持更大批次(如单卡batch=4,全局batch=128)。
- 应用场景:大规模图像分类(如ImageNet)。
- 案例:单机8卡训练ViT-Large模型。
torchrun --nproc_per_node=8 --master_port=12345 train.py
- 优势:线性加速比,训练时间从单卡7天缩短至多卡1天。
- 应用场景:训练多模态感知模型(如激光雷达+摄像头融合)。
- 案例:在2机16卡集群中训练3D物体检测模型。
# 主节点(IP: 10.1.1.1) torchrun --nnodes=2 --nproc_per_node=8 --node_rank=0 --master_addr=10.1.1.1 train.py # 从节点(IP: 10.1.1.2) torchrun --nnodes=2 --nproc_per_node=8 --node_rank=1 --master_addr=10.1.1.1 train.py
- 优势:处理高分辨率点云数据(单卡显存不足时,分布式训练分割数据)。
- 核心价值:
torchrun简化了分布式训练的部署,支持灵活的多机多卡配置。 - 领域适配性:
- AIGC:加速生成模型训练,支撑创意快速迭代。
- 传统深度学习:降低大规模数据训练成本。
- 自动驾驶:解决高算力需求场景的显存与计算瓶颈。
- 面试加分点:结合业务场景(如自动驾驶需低延迟)说明显存优化策略(如梯度累积)。



