PyTorch 入门

import torch
import numpy as np

# 创建张量
a = torch.tensor([1, 2, 3])
b = torch.zeros(3, 4)           # 全 0
c = torch.ones(2, 3)            # 全 1
d = torch.randn(3, 4)           # 标准正态分布
e = torch.arange(0, 10, 2)     # [0, 2, 4, 6, 8]

# 从 NumPy 互转
arr = np.array([1, 2, 3])
t = torch.from_numpy(arr)       # ndarray → tensor
arr2 = t.numpy()                # tensor → ndarray（仅 CPU 上有效）

# 属性
print(d.shape)      # torch.Size([3, 4])
print(d.dtype)      # torch.float32
print(d.device)     # cpu

# GPU 迁移
if torch.cuda.is_available():
    d = d.cuda()    # 或 d.to('cuda')

x = torch.randn(3, 4)
y = torch.randn(3, 4)

# 逐元素运算
z = x + y        # 加法
z = x * y        # 逐元素乘法（不是矩阵乘法）
z = torch.exp(x)

# 矩阵乘法
z = x @ y.T      # @ 是矩阵乘法
z = torch.mm(x, y.T)   # 等价写法

# 维度操作
x = torch.randn(2, 3, 4)
x = x.reshape(2, 12)        # 变形
x = x.view(-1, 12)          # 类似 reshape，但要求内存连续
x = x.squeeze()             # 去掉大小为 1 的维度
x = x.unsqueeze(0)          # 在位置 0 添加一个维度

# 索引
x[:, 0, :]                  # 类似 NumPy 切片
x[x > 0]                    # 布尔索引
torch.cat([x, y], dim=0)    # 拼接

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

class SimpleNet(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, output_dim)
        self.dropout = nn.Dropout(0.3)

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

model = SimpleNet(784, 256, 10)
print(model)
# 打印参数量
print(f"参数量: {sum(p.numel() for p in model.parameters()):,}")

from torch.utils.data import DataLoader, Dataset

class MyDataset(Dataset):
    def __init__(self, data, labels):
        self.data = data
        self.labels = labels

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

    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]

# 通常使用 torchvision 内置数据集
from torchvision import datasets, transforms

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

train_data = datasets.MNIST(
    root='./data', train=True,
    download=True, transform=transform
)

train_loader = DataLoader(
    train_data, batch_size=64,
    shuffle=True, num_workers=4
)

import torch.optim as optim

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = SimpleNet(784, 256, 10).to(device)

criterion = nn.CrossEntropyLoss()    # 损失函数
optimizer = optim.Adam(model.parameters(), lr=0.001)

num_epochs = 10
for epoch in range(num_epochs):
    model.train()
    total_loss = 0
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        data = data.view(data.size(0), -1)  # 展平图片

        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

        total_loss += loss.item()

    print(f"Epoch {epoch+1}, Loss: {total_loss/len(train_loader):.4f}")

# 维度不匹配：打印每层 shape
print(x.shape)

# GPU 显存溢出
torch.cuda.empty_cache()
# 减小 batch_size 或用梯度累积

# 梯度爆炸/消失
# 打印梯度范数
total_norm = 0
for p in model.parameters():
    if p.grad is not None:
        total_norm += p.grad.norm().item() ** 2
print(f"Gradient norm: {total_norm**0.5:.4f}")
# > 10 可能爆炸，< 1e-6 可能消失

# 过拟合
# 加 Dropout、数据增强、weight decay、早停