深度学习框架PyTorch实战

前言

PyTorch作为最受欢迎的深度学习框架之一，以其灵活性和易用性赢得了研究者和工程师的青睐。本文将带你从零开始学习PyTorch，通过实战案例掌握深度学习的核心技能。

一、PyTorch基础

1.1 安装PyTorch

1
2
3
4
5
6
7
8
# CPU版本
pip install torch torchvision torchaudio

# GPU版本（CUDA 11.8）
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

# 验证安装
python -c "import torch; print(torch.__version__); print(torch.cuda.is_available())"

1.2 张量（Tensor）基础

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
import torch
import numpy as np

# 创建张量
x = torch.tensor([1, 2, 3])
print(x)  # tensor([1, 2, 3])

# 从numpy创建
arr = np.array([1, 2, 3])
x = torch.from_numpy(arr)

# 创建特殊张量
zeros = torch.zeros(2, 3)          # 全0
ones = torch.ones(2, 3)             # 全1
random = torch.randn(2, 3)          # 标准正态分布
uniform = torch.rand(2, 3)          # [0, 1)均匀分布
arange = torch.arange(0, 10, 2)     # 等差数列

# 张量属性
print(x.shape)        # torch.Size([3])
print(x.dtype)        # torch.int64
print(x.device)       # cpu

1.3 张量操作

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
# 数学运算
a = torch.tensor([1.0, 2.0, 3.0])
b = torch.tensor([4.0, 5.0, 6.0])

c = a + b              # 加法
c = torch.add(a, b)    # 等价写法
c = a * b              # 元素wise乘法
c = torch.matmul(a, b) # 矩阵乘法（对于向量是点积）

# 形状操作
x = torch.randn(2, 3, 4)
y = x.view(2, 12)        # 改变形状
z = x.reshape(2, -1)     # -1自动推导
w = x.transpose(1, 2)    # 转置

# 索引和切片
x = torch.randn(3, 4)
print(x[0, :])           # 第一行
print(x[:, 1])           # 第二列
print(x[0:2, 1:3])       # 切片

# 拼接
x = torch.randn(2, 3)
y = torch.randn(2, 3)
z = torch.cat([x, y], dim=0)  # 按行拼接 (4, 3)
z = torch.cat([x, y], dim=1)  # 按列拼接 (2, 6)
z = torch.stack([x, y], dim=0) # 堆叠 (2, 2, 3)

1.4 GPU加速

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
# 检查GPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'Using device: {device}')

# 将张量移到GPU
x = torch.randn(1000, 1000)
x = x.to(device)
# 或
x = x.cuda()  # 移到GPU
x = x.cpu()   # 移回CPU

# 创建时指定设备
x = torch.randn(1000, 1000, device=device)

# GPU上的计算
a = torch.randn(1000, 1000, device='cuda')
b = torch.randn(1000, 1000, device='cuda')
c = torch.matmul(a, b)  # GPU上执行

二、自动微分

2.1 基础自动微分

1
2
3
4
5
6
7
8
9
10
11
12
# 创建需要梯度的张量
x = torch.tensor([2.0], requires_grad=True)
y = x ** 2 + 3 * x + 1

# 反向传播
y.backward()

# 查看梯度
print(x.grad)  # dy/dx = 2x + 3 = 7

# 梯度清零
x.grad.zero_()

2.2 复杂计算图

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
# 多变量函数
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
w = torch.tensor([0.5, 0.3, 0.2], requires_grad=True)
b = torch.tensor([1.0], requires_grad=True)

# 前向传播
y = torch.sum(x * w) + b
loss = y ** 2

# 反向传播
loss.backward()

print(f'x.grad: {x.grad}')  # dL/dx
print(f'w.grad: {w.grad}')  # dL/dw
print(f'b.grad: {b.grad}')  # dL/db

2.3 梯度管理

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
# 禁用梯度计算（推理时）
with torch.no_grad():
    y = x * 2
    print(y.requires_grad)  # False

# 或使用装饰器
@torch.no_grad()
def inference(x):
    return x * 2

# 分离计算图
x = torch.tensor([1.0], requires_grad=True)
y = x ** 2
z = y.detach()  # z不需要梯度

# 临时启用梯度
with torch.enable_grad():
    y = x * 2

三、构建神经网络

3.1 nn.Module基础

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
import torch.nn as nn
import torch.nn.functional as F

# 定义神经网络
class SimpleNN(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, output_size)
    
    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 创建模型
model = SimpleNN(784, 128, 10)

# 查看模型结构
print(model)

# 查看参数
for name, param in model.named_parameters():
    print(f'{name}: {param.shape}')

# 前向传播
x = torch.randn(32, 784)  # batch_size=32
output = model(x)
print(output.shape)  # torch.Size([32, 10])

3.2 常用层

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
# 全连接层
fc = nn.Linear(in_features=128, out_features=64)

# 卷积层
conv2d = nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, padding=1)

# 池化层
maxpool = nn.MaxPool2d(kernel_size=2, stride=2)
avgpool = nn.AvgPool2d(kernel_size=2)

# Dropout
dropout = nn.Dropout(p=0.5)

# 批归一化
bn = nn.BatchNorm2d(num_features=64)

# 激活函数
relu = nn.ReLU()
sigmoid = nn.Sigmoid()
tanh = nn.Tanh()

3.3 Sequential容器

1
2
3
4
5
6
7
8
9
10
11
12
13
14
# 使用Sequential简化定义
model = nn.Sequential(
    nn.Linear(784, 256),
    nn.ReLU(),
    nn.Dropout(0.2),
    nn.Linear(256, 128),
    nn.ReLU(),
    nn.Dropout(0.2),
    nn.Linear(128, 10)
)

# 前向传播
x = torch.randn(32, 784)
output = model(x)

3.4 CNN示例

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
class CNN(nn.Module):
    def __init__(self, num_classes=10):
        super(CNN, self).__init__()
        
        # 卷积层
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        
        # 池化层
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        
        # 全连接层
        self.fc1 = nn.Linear(64 * 7 * 7, 128)
        self.fc2 = nn.Linear(128, num_classes)
        
        # Dropout
        self.dropout = nn.Dropout(0.5)
    
    def forward(self, x):
        # x: (batch, 1, 28, 28)
        
        x = F.relu(self.conv1(x))   # (batch, 32, 28, 28)
        x = self.pool(x)             # (batch, 32, 14, 14)
        
        x = F.relu(self.conv2(x))   # (batch, 64, 14, 14)
        x = self.pool(x)             # (batch, 64, 7, 7)
        
        x = x.view(x.size(0), -1)    # (batch, 64*7*7)
        
        x = F.relu(self.fc1(x))      # (batch, 128)
        x = self.dropout(x)
        x = self.fc2(x)              # (batch, 10)
        
        return x

model = CNN(num_classes=10)

四、模型训练

4.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
52
53
54
55
56
57
58
59
60
61
62
63
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

# 1. 准备数据
X_train = torch.randn(1000, 784)
y_train = torch.randint(0, 10, (1000,))

train_dataset = TensorDataset(X_train, y_train)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

# 2. 定义模型
model = SimpleNN(784, 128, 10)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)

# 3. 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 4. 训练循环
num_epochs = 10

for epoch in range(num_epochs):
    model.train()  # 训练模式
    running_loss = 0.0
    correct = 0
    total = 0
    
    for batch_idx, (data, target) in enumerate(train_loader):
        # 移到GPU
        data, target = data.to(device), target.to(device)
        
        # 梯度清零
        optimizer.zero_grad()
        
        # 前向传播
        outputs = model(data)
        loss = criterion(outputs, target)
        
        # 反向传播
        loss.backward()
        
        # 更新参数
        optimizer.step()
        
        # 统计
        running_loss += loss.item()
        _, predicted = torch.max(outputs.data, 1)
        total += target.size(0)
        correct += (predicted == target).sum().item()
        
        if (batch_idx + 1) % 10 == 0:
            print(f'Epoch [{epoch+1}/{num_epochs}], '
                  f'Step [{batch_idx+1}/{len(train_loader)}], '
                  f'Loss: {loss.item():.4f}')
    
    # 每个epoch的统计
    epoch_loss = running_loss / len(train_loader)
    epoch_acc = 100 * correct / total
    print(f'Epoch [{epoch+1}/{num_epochs}] - '
          f'Loss: {epoch_loss:.4f}, Accuracy: {epoch_acc:.2f}%')

4.2 验证和测试

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
def evaluate(model, data_loader, device):
    model.eval()  # 评估模式
    correct = 0
    total = 0
    
    with torch.no_grad():
        for data, target in data_loader:
            data, target = data.to(device), target.to(device)
            outputs = model(data)
            _, predicted = torch.max(outputs.data, 1)
            total += target.size(0)
            correct += (predicted == target).sum().item()
    
    accuracy = 100 * correct / total
    return accuracy

# 在测试集上评估
test_acc = evaluate(model, test_loader, device)
print(f'Test Accuracy: {test_acc:.2f}%')

4.3 学习率调度

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
# StepLR：每N个epoch衰减
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)

# ReduceLROnPlateau：根据指标自动调整
scheduler = optim.lr_scheduler.ReduceLROnPlateau(
    optimizer, mode='min', factor=0.5, patience=5
)

# CosineAnnealingLR：余弦退火
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10)

# 训练循环中使用
for epoch in range(num_epochs):
    train(...)
    val_loss = validate(...)
    
    # 更新学习率
    scheduler.step()  # StepLR
    # 或
    scheduler.step(val_loss)  # ReduceLROnPlateau

五、实战：MNIST手写数字识别

5.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
from torchvision import datasets, transforms

# 数据预处理
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

# 下载数据
train_dataset = datasets.MNIST(
    root='./data',
    train=True,
    download=True,
    transform=transform
)

test_dataset = datasets.MNIST(
    root='./data',
    train=False,
    transform=transform
)

# 创建DataLoader
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False)

5.2 定义模型

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
class MNISTNet(nn.Module):
    def __init__(self):
        super(MNISTNet, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

model = MNISTNet().to(device)

5.3 训练和评估

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
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.NLLLoss()

def train(model, device, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        
        if batch_idx % 100 == 0:
            print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)}] '
                  f'Loss: {loss.item():.6f}')

def test(model, device, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += criterion(output, target).item()
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)
    accuracy = 100. * correct / len(test_loader.dataset)
    print(f'\nTest set: Average loss: {test_loss:.4f}, '
          f'Accuracy: {correct}/{len(test_loader.dataset)} ({accuracy:.2f}%)\n')

# 训练
for epoch in range(1, 11):
    train(model, device, train_loader, optimizer, epoch)
    test(model, device, test_loader)

六、模型保存与加载

6.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
# 保存整个模型
torch.save(model, 'model.pth')

# 加载模型
model = torch.load('model.pth')

# 保存模型参数（推荐）
torch.save(model.state_dict(), 'model_weights.pth')

# 加载参数
model = MNISTNet()
model.load_state_dict(torch.load('model_weights.pth'))
model.eval()

# 保存训练状态
checkpoint = {
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': loss,
}
torch.save(checkpoint, 'checkpoint.pth')

# 恢复训练
checkpoint = torch.load('checkpoint.pth')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']

七、总结

本文介绍了PyTorch的核心概念和实战应用：

关键知识点

  1. 张量操作：PyTorch的基础数据结构
  2. 自动微分：反向传播的核心
  3. nn.Module：构建神经网络的基石
  4. 训练循环：完整的模型训练流程

下一步学习

  • 迁移学习（Transfer Learning）
  • 对抗生成网络（GAN）
  • Transformer架构
  • 模型部署（TorchScript、ONNX）

掌握PyTorch，开启你的深度学习之旅！

参考资料