简介:本文为PyTorch开发者提供从基础安装到高级技巧的全流程指导,涵盖张量操作、模型构建、分布式训练等核心模块,结合代码示例与实战建议,助力快速掌握深度学习框架。
PyTorch提供CPU与GPU(CUDA)双版本,推荐通过官方命令安装最新稳定版:
# 示例:安装支持CUDA 11.8的PyTorch 2.0
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
需注意:
nvidia-smi
查看驱动支持的CUDA版本,选择对应PyTorch版本conda create -n pytorch_env python=3.9
),避免依赖冲突运行以下代码验证环境:
import torch
print(torch.__version__) # 输出版本号
print(torch.cuda.is_available()) # 输出True表示GPU可用
PyTorch张量支持与NumPy数组的无缝转换:
import torch
import numpy as np
# 从列表创建
x = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32)
# 从NumPy数组转换
np_array = np.random.rand(2, 2)
y = torch.from_numpy(np_array)
# 张量运算
z = x * 2 + torch.sin(x) # 支持广播机制
通过requires_grad=True
启用梯度计算:
x = torch.tensor(2.0, requires_grad=True)
y = x ** 3 + 5 * x
y.backward() # 计算dy/dx
print(x.grad) # 输出梯度值:3*2^2 + 5 = 17
关键点:
optimizer.zero_grad()
with torch.no_grad():
禁用梯度计算以节省内存自定义模型需继承nn.Module
并实现forward()
方法:
import torch.nn as nn
class CNN(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(3, 16, kernel_size=3)
self.fc = nn.Linear(16*14*14, 10)
def forward(self, x):
x = torch.relu(self.conv1(x))
x = x.view(x.size(0), -1) # 展平
return self.fc(x)
常用组合示例:
model = CNN()
criterion = nn.CrossEntropyLoss() # 分类任务
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
# 训练步骤
outputs = model(inputs)
loss = criterion(outputs, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
自定义数据集类:
from torch.utils.data import Dataset, DataLoader
class CustomDataset(Dataset):
def __init__(self, data, labels):
self.data = data
self.labels = labels
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
return self.data[idx], self.labels[idx]
# 使用示例
dataset = CustomDataset(train_data, train_labels)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)
通过torchvision.transforms
实现:
from torchvision import transforms
transform = transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
使用DataParallel
快速实现多卡并行:
if torch.cuda.device_count() > 1:
model = nn.DataParallel(model)
model.to('cuda')
或使用更高效的DistributedDataParallel
(DDP):
import torch.distributed as dist
dist.init_process_group('nccl')
model = nn.parallel.DistributedDataParallel(model)
通过torch.cuda.amp
减少显存占用:
scaler = torch.cuda.amp.GradScaler()
with torch.cuda.amp.autocast():
outputs = model(inputs)
loss = criterion(outputs, labels)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
将PyTorch模型转换为可序列化格式:
# 跟踪模式(适用于动态图)
traced_script = torch.jit.trace(model, example_input)
traced_script.save("model.pt")
# 脚本模式(适用于控制流)
scripted_model = torch.jit.script(model)
与TensorFlow等框架互操作:
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(
model,
dummy_input,
"model.onnx",
input_names=["input"],
output_names=["output"]
)
调试工具:
torch.autograd.set_detect_anomaly(True)
捕获梯度异常torchsummary
可视化模型结构:
from torchsummary import summary
summary(model, input_size=(3, 224, 224))
性能优化:
torch.nn.functional
中的函数式接口torch.no_grad()
减少计算图构建常见问题:
batch_size
或使用梯度累积本手册系统梳理了PyTorch从基础到进阶的核心知识,结合代码示例与工程实践建议,适合不同阶段的开发者快速掌握深度学习框架的核心能力。建议读者结合官方文档(pytorch.org/docs)持续学习最新特性。