PyTorch深度学习框架详解

PyTorch是当前最流行的深度学习框架之一，以动态计算图和Pythonic API著称，广泛应用于研究和生产环境。

张量基础

张量是PyTorch的基本数据结构，类似于NumPy数组但支持GPU加速。张量操作是深度学习计算的基础。

import torch

# 创建张量
x = torch.tensor([1.0, 2.0, 3.0])
zeros = torch.zeros(3, 4)
ones = torch.ones(2, 3)
rand = torch.randn(3, 3)

# 基本运算
a = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
b = torch.tensor([[5.0, 6.0], [7.0, 8.0]])

c = a + b
d = torch.mm(a, b)  # 矩阵乘法
e = torch.dot(a.flatten(), b.flatten())  # 点积

print(f"张量形状: {a.shape}")
print(f"矩阵乘法结果:\n{d}")

autograd自动微分

autograd是PyTorch的自动微分引擎，通过构建动态计算图自动计算梯度，支持反向传播训练。

# 启用梯度计算
x = torch.tensor(2.0, requires_grad=True)
w = torch.tensor(3.0, requires_grad=True)
b = torch.tensor(1.0, requires_grad=True)

# 前向传播
y = w * x + b  # y = 3*2 + 1 = 7
z = y ** 2     # z = 49

# 反向传播
z.backward()

print(f"dz/dw = {w.grad}")  # dz/dw = 2*y*x = 28
print(f"dz/dx = {x.grad}")  # dz/dx = 2*y*w = 42
print(f"dz/db = {b.grad}")  # dz/db = 2*y = 14

nn.Module模型构建

nn.Module是构建神经网络的基类，封装参数管理和前向传播逻辑，支持模型序列化和设备迁移。

import torch.nn as nn

class CustomModel(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        self.layer1 = nn.Linear(input_dim, hidden_dim)
        self.layer2 = nn.Linear(hidden_dim, hidden_dim)
        self.layer3 = nn.Linear(hidden_dim, output_dim)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.2)
    
    def forward(self, x):
        x = self.relu(self.layer1(x))
        x = self.dropout(x)
        x = self.relu(self.layer2(x))
        x = self.layer3(x)
        return x

model = CustomModel(784, 256, 10)
print(f"模型参数量: {sum(p.numel() for p in model.parameters())}")

训练循环

完整的训练循环包括数据加载、前向传播、损失计算、反向传播和参数更新五个步骤。

import torch.optim as optim

def train_model(model, train_loader, epochs=10, lr=0.001):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = model.to(device)
    
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=lr)
    
    for epoch in range(epochs):
        model.train()
        total_loss = 0.0
        
        for batch_x, batch_y in train_loader:
            batch_x, batch_y = batch_x.to(device), batch_y.to(device)
            
            optimizer.zero_grad()
            outputs = model(batch_x)
            loss = criterion(outputs, batch_y)
            loss.backward()
            optimizer.step()
            
            total_loss += loss.item()
        
        avg_loss = total_loss / len(train_loader)
        print(f"Epoch {epoch+1}/{epochs}, Loss: {avg_loss:.4f}")

# 示例训练数据
X_train = torch.randn(1000, 784)
y_train = torch.randint(0, 10, (1000,))
dataset = torch.utils.data.TensorDataset(X_train, y_train)
loader = torch.utils.data.DataLoader(dataset, batch_size=32, shuffle=True)

模型保存与加载

PyTorch支持保存和加载模型参数或整个模型，便于模型持久化和迁移部署。

# 保存模型参数
torch.save(model.state_dict(), 'model.pth')

# 加载模型参数
loaded_model = CustomModel(784, 256, 10)
loaded_model.load_state_dict(torch.load('model.pth'))
loaded_model.eval()  # 设置为评估模式

print("模型已保存和加载完成")

PyTorch以其灵活性和易用性成为深度学习研究的首选框架，掌握这些基础组件是构建复杂模型的关键。