← 返回首页
🤖

PyTorch基础教程

📂 ai ⏱ 2 min 273 words

PyTorch基础教程

PyTorch简介

PyTorch是Facebook开发的深度学习框架,以动态计算图和Pythonic设计著称,广泛用于研究和生产。

张量操作

import torch

# 创建张量
a = torch.zeros(3, 4)
b = torch.ones(3, 4)
c = torch.randn(3, 4)
d = torch.tensor([1, 2, 3], dtype=torch.float32)

print(f"零张量: {a.shape}")
print(f"随机张量: {c.shape}")

张量运算

# 基本运算
x = torch.randn(2, 3)
y = torch.randn(2, 3)

add = x + y           # 逐元素加法
mul = x * y           # 逐元素乘法
matmul = torch.mm(x, y.t())  # 矩阵乘法

print(f"矩阵乘法结果: {matmul.shape}")

# 形状操作
reshape = torch.randn(2, 6).view(2, 3, 2)
print(f"重塑后: {reshape.shape}")

# 聚合操作
x = torch.randn(3, 4)
print(f"求和: {x.sum(dim=0).shape}")
print(f"均值: {x.mean(dim=1).shape}")

自动微分

# requires_grad启用梯度计算
x = torch.tensor([2.0, 3.0], requires_grad=True)
y = x ** 2 + 2 * x + 1

# 反向传播
y.sum().backward()
print(f"梯度 dy/dx = 2x+2: {x.grad}")  # [6.0, 8.0]

# 阻止梯度计算
with torch.no_grad():
    y_no_grad = x ** 2

nn.Module

import torch.nn as nn

class CustomModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(CustomModel, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, output_size)
    
    def forward(self, x):
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return x

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

完整训练流程

from torch.utils.data import DataLoader, TensorDataset

# 1. 准备数据
X = torch.randn(500, 10)
y = torch.randint(0, 2, (500,))
dataset = TensorDataset(X, y)
loader = DataLoader(dataset, batch_size=32, shuffle=True)

# 2. 创建模型
model = CustomModel(10, 20, 2)

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

# 4. 训练循环
for epoch in range(10):
    model.train()
    total_loss = 0
    for batch_x, batch_y in loader:
        optimizer.zero_grad()
        output = model(batch_x)
        loss = criterion(output, batch_y)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    
    if (epoch + 1) % 5 == 0:
        print(f"Epoch [{epoch+1}/10], Loss: {total_loss/len(loader):.4f}")

# 5. 评估
model.eval()
with torch.no_grad():
    output = model(X)
    predictions = output.argmax(dim=1)
    accuracy = (predictions == y).float().mean()
    print(f"准确率: {accuracy:.4f}")

保存与加载模型

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

# 加载模型
loaded_model = CustomModel(10, 20, 2)
loaded_model.load_state_dict(torch.load('model.pth'))

总结

PyTorch以直观的API和动态图机制,成为深度学习研究的首选框架之一。