PyTorch是一个广泛使用的深度学习框架,它提供了许多功能强大的工具和函数,用于构建和训练神经网络。以下是一些PyTorch的常用知识点和示例说明:
-
张量(Tensors)
-
创建张量:使用torch.tensor()、torch.Tensor()或特定创建函数如torch.zeros(), torch.ones(), torch.randn()等创建不同类型的张量。
import torch x = torch.tensor([1., 2., 3.]) # 创建一个浮点型张量 zeros_tensor = torch.zeros((3, 4)) # 创建一个3x4的全零张量
-
张量操作:类似NumPy,支持各种数学运算和索引操作,如加减乘除、矩阵乘法、广播机制、切片等。
y = torch.tensor([4., 5., 6.]) result = x + y # 张量加法
-
数据类型转换:通过.to()方法可以改变张量的数据类型或者设备(CPU/GPU)。
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') x_gpu = x.to(device) # 将张量移动到GPU上
-
-
自动微分(Autograd)
-
使用.requires_grad_()标记张量以启用梯度计算:
x.requires_grad_() y = x * 2 z = y.sum() z.backward() # 自动计算梯度 print(x.grad) # 输出x的梯度
-
-
神经网络模块(nn.Module)
-
定义网络结构:继承自nn.Module并实现__init__和forward方法。
import torch.nn as nn class SimpleNet(nn.Module): def __init__(self): super(SimpleNet, self).__init__() self.linear = nn.Linear(784, 10) # 定义一个线性层 def forward(self, x): out = self.linear(x) return out
-
构建与训练模型:
model = SimpleNet() criterion = nn.CrossEntropyLoss() optimizer = torch.optim.SGD(model.parameters(), lr=0.01) for epoch in range(num_epochs): # 前向传播 output = model(inputs) loss = criterion(output, targets) # 反向传播及优化 optimizer.zero_grad() loss.backward() optimizer.step()
-
-
数据加载器(DataLoader)
-
使用torch.utils.data.DataLoader来加载和批处理数据。
from torch.utils.data import DataLoader, TensorDataset dataset = TensorDataset(data_tensor, label_tensor) dataloader = DataLoader(dataset, batch_size=64, shuffle=True) for batch_data, batch_labels in dataloader: # 在每个迭代周期中,batch_data和batch_labels会是当前批次的张量数据 pass
-
-
保存与加载模型
-
使用torch.save()和torch.load()保存和加载模型参数或整个模型。
torch.save(model.state_dict(), 'model.pth') # 保存模型参数 model.load_state_dict(torch.load('model.pth')) # 加载模型参数 # 或者保存整个模型 torch.save(model, 'model_full.pth') # 保存整个模型(包括其结构和参数) loaded_model = torch.load('model_full.pth', map_location=device) # 加载整个模型
-
-
多GPU并行训练
-
使用torch.nn.DataParallel或torch.nn.parallel.DistributedDataParallel进行多GPU训练。
model = nn.DataParallel(SimpleNet()) # 如果有多块GPU可用,则将模型分布到多个GPU上
-
-
控制流(autograd with control flow)
- PyTorch支持在动态图模式下使用Python的控制流语句(如if-else、for循环),并且能正确跟踪梯度。
动态计算图、混合精度训练、量化压缩、可视化工具
动态计算图(Dynamic Computation Graph)
在PyTorch中,计算图是在运行时构建的,这意味着你可以根据程序运行的状态实时改变网络结构或执行不同的计算路径。这是与静态计算图框架如TensorFlow的一个显著区别。
示例:
# 动态改变模型结构
class DynamicModel(nn.Module):
def __init__(self):
super(DynamicModel, self).__init__()
self.linear1 = nn.Linear(10, 5)
self.linear2 = nn.Linear(5, 3)
def forward(self, x, use_second_layer=True):
out = F.relu(self.linear1(x))
if use_second_layer:
out = self.linear2(out) # 根据条件决定是否使用第二层
return out
model = DynamicModel()
混合精度训练(Mixed Precision Training)
混合精度训练利用了FP16和FP32数据类型的优势,通过将部分计算转移到半精度上以减少内存占用和加快计算速度,同时保持关键部分(如梯度更新)在全精度下进行,以维持数值稳定性。
使用torch.cuda.amp模块实现自动混合精度训练:
import torch
from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()
for inputs, targets in dataloader:
inputs = inputs.cuda()
targets = targets.cuda()
with autocast():
outputs = model(inputs)
loss = criterion(outputs, targets)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad()
量化压缩(Quantization)
量化是将模型的权重和激活从浮点数转换为低比特整数的过程,从而减小模型大小并加速推理。PyTorch提供了量化API来实现这一过程。
简化版量化示例:
import torch.quantization
# 假设model是一个已经训练好的模型
model_fp32 = ... # 初始化并训练模型
# 首先对模型进行伪量化(模拟量化)
prepared_model = torch.quantization.prepare(model_fp32)
# 进行量化校准(收集统计数据)
quantized_model = torch.quantization.convert(prepared_model)
# 现在quantized_model是一个量化后的模型,可以用于推理
可视化工具
PyTorch支持通过torchviz库来进行计算图可视化,或者配合其他工具(如TensorBoard)展示模型结构、训练指标等。
对于简单的计算图可视化:
from torchviz import make_dot
x = torch.randn(5, requires_grad=True)
y = x * 2
z = y ** 2
z.backward(torch.ones_like(z))
dot_graph = make_dot(z)
dot_graph.view() # 在Jupyter Notebook中显示图形
对于模型结构可视化,通常结合torchsummary或直接使用TensorBoard配合torch.utils.tensorboard接口:
from torchsummary import summary
summary(model, input_size=(1, 28, 28)) # 对于卷积神经网络,输入维度为(通道, 高, 宽)
# 或者在TensorBoard中展示模型结构
from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter()
writer.add_graph(model, torch.rand((1, 28, 28))) # 输入一个随机张量获取模型结构
writer.close()