速查手册

import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self, in_dim, hidden_dim, out_dim):
        super().__init__()
        self.fc1 = nn.Linear(in_dim, hidden_dim)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.5)
        self.fc2 = nn.Linear(hidden_dim, out_dim)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.dropout(x)
        x = self.fc2(x)
        return x

model = MyModel(784, 256, 10)
print(sum(p.numel() for p in model.parameters()))  # 参数量

model = MyModel(784, 256, 10).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(num_epochs):
    model.train()                       # 训练模式（启用 Dropout/BN）
    for x, y in train_loader:
        x, y = x.to(device), y.to(device)
        optimizer.zero_grad()           # 清空梯度
        pred = model(x)                 # 前向传播
        loss = criterion(pred, y)       # 计算损失
        loss.backward()                 # 反向传播
        optimizer.step()                # 更新参数

    # 验证
    model.eval()                        # 推理模式
    with torch.no_grad():               # 不计算梯度
        correct = 0
        for x, y in val_loader:
            pred = model(x.to(device))
            correct += (pred.argmax(1) == y.to(device)).sum().item()
    print(f"Epoch {epoch}: val_acc={correct/val_total:.4f}")

from torch.utils.data import Dataset, DataLoader
from torchvision import transforms

# 自定义 Dataset
class MyDataset(Dataset):
    def __init__(self, data, labels, transform=None):
        self.data = data
        self.labels = labels
        self.transform = transform

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        x, y = self.data[idx], self.labels[idx]
        if self.transform:
            x = self.transform(x)
        return x, y

# 常用 transforms
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5], std=[0.5]),
    transforms.RandomHorizontalFlip(),   # 数据增强
])

# DataLoader
loader = DataLoader(
    MyDataset(data, labels, transform),
    batch_size=64, shuffle=True, num_workers=4
)

# 模型保存与加载
torch.save(model.state_dict(), 'model.pth')
model.load_state_dict(torch.load('model.pth'))