PyTorch 深度学习框架入门指南

5.2 基本运算

a = torch.tensor([1.0, 2.0])
b = torch.tensor([3.0, 4.0])

# 加法
c = a + b

# 矩阵乘法
d = torch.matmul(a, b)

# 查看设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
a = a.to(device)

6. 构建简单神经网络

PyTorch 中定义模型通常继承自 nn.Module 类。

import torch.nn as nn
import torch.nn.functional as F

class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        # 定义全连接层
        self.fc1 = nn.Linear(784, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        # 展平输入
        x = x.view(-1, 784)
        # 第一层 + ReLU
        x = F.relu(self.fc1(x))
        # 第二层
        x = self.fc2(x)
        return x

model = SimpleNet()
print(model)

7. 训练流程详解

一个完整的训练循环包含数据加载、前向传播、损失计算、反向传播和优化器更新。

import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

# 假设已有数据 X_train, y_train
X_train = torch.randn(100, 784)
y_train = torch.randint(0, 10, (100,))

dataset = TensorDataset(X_train, y_train)
dataloader = DataLoader(dataset, batch_size=10, shuffle=True)

model = SimpleNet()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 训练循环
for epoch in range(10):
    for batch_X, batch_y in dataloader:
        # 前向传播
        outputs = model(batch_X)
        loss = criterion(outputs, batch_y)
        
        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}")

8. 总结

PyTorch 以其灵活的动态图和简洁的 API 成为深度学习开发的首选框架之一。掌握张量操作、自动求导机制以及模型构建流程，是深入使用 PyTorch 的基础。通过官方文档和持续实践，开发者可以高效地构建和部署各类深度学习模型。