PyTorch 完全征服指南：从入门到精通 — 从张量到分布式训练

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu

import torch

# 检查 PyTorch 版本
print(f"PyTorch 版本: {torch.__version__}")

# 检查 CUDA 是否可用
print(f"CUDA 可用: {torch.cuda.is_available()}")

# 检查 GPU 数量
if torch.cuda.is_available():
    print(f"GPU 数量: {torch.cuda.device_count()}")
    print(f"当前 GPU: {torch.cuda.get_device_name(0)}")
    print(f"GPU 显存: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")

# 检查 Apple Silicon (M1/M2/M3) 的 MPS
print(f"MPS 可用: {torch.backends.mps.is_available()}")

# 自动选择设备
if torch.cuda.is_available():
    device = torch.device("cuda")
elif torch.backends.mps.is_available():
    device = torch.device("mps")
else:
    device = torch.device("cpu")

print(f"使用设备: {device}")

import torch
import numpy as np

# 直接从数据创建
t1 = torch.tensor([1, 2, 3, 4, 5])
print(f"1D 张量: {t1}, shape: {t1.shape}, dtype: {t1.dtype}")

# 2D 张量（矩阵）
t2 = torch.tensor([[1.0, 2.0, 3.0],
                   [4.0, 5.0, 6.0]])
print(f"2D 张量:\n{t2}, shape: {t2.shape}")

# 特殊张量的创建
zeros = torch.zeros(3, 4)          # 全为 0
ones = torch.ones(2, 3)            # 全为 1
rand = torch.rand(3, 3)            # 0~1 均匀分布
randn = torch.randn(3, 3)          # 标准正态分布
eye = torch.eye(4)                  # 单位矩阵
arange = torch.arange(0, 10, 2)    # [0, 2, 4, 6, 8]
linspace = torch.linspace(0, 1, 5) # 均匀间隔 5 个值

print(f"zeros:\n{zeros}")
print(f"randn:\n{randn}")

# 创建与已有张量形状相同的张量
t3 = torch.zeros_like(t2)
t4 = torch.ones_like(t2)
t5 = torch.rand_like(t2)

# 从 NumPy 数组创建（共享内存）
np_arr = np.array([1.0, 2.0, 3.0])
t_from_np = torch.from_numpy(np_arr)
print(f"来自 NumPy: {t_from_np}")

# 张量转换为 NumPy（仅限 CPU）
np_from_t = t1.numpy()

t = torch.rand(3, 4, 5)

# 基本属性
print(f"shape: {t.shape}")       # torch.Size([3, 4, 5])
print(f"ndim: {t.ndim}")         # 3
print(f"dtype: {t.dtype}")       # torch.float32
print(f"device: {t.device}")     # cpu
print(f"numel: {t.numel()}")     # 60（元素总数）

# 数据类型转换
t_int = t.to(torch.int32)
t_long = t.long()           # torch.int64
t_float = t.float()         # torch.float32
t_double = t.double()       # torch.float64
t_half = t.half()           # torch.float16

# 移动到 GPU
if torch.cuda.is_available():
    t_gpu = t.to("cuda")
    t_gpu2 = t.cuda()        # 结果相同
    t_back = t_gpu.cpu()     # 再移回 CPU

t = torch.arange(24)  # 0~23 的 1D 张量

# reshape：只要元素数相同，可以变换为任意形状
t_2d = t.reshape(4, 6)
t_3d = t.reshape(2, 3, 4)
t_auto = t.reshape(6, -1)  # -1 表示自动计算（6x4）

# view：与 reshape 类似，但要求内存连续
t_view = t.view(3, 8)

# squeeze/unsqueeze：移除/增加维度
t = torch.zeros(1, 3, 1, 4)
print(f"原始 shape: {t.shape}")  # [1, 3, 1, 4]

t_sq = t.squeeze()       # 移除大小为 1 的维度 → [3, 4]
t_sq1 = t.squeeze(0)     # 只移除第 0 维 → [3, 1, 4]
t_unsq = t_sq.unsqueeze(0)  # 在第 0 个位置增加维度 → [1, 3, 4]

# transpose/permute：改变维度顺序
t = torch.rand(2, 3, 4)
t_T = t.transpose(0, 1)    # [3, 2, 4]
t_perm = t.permute(2, 0, 1)  # [4, 2, 3]

# contiguous：permute 之后保证内存连续
t_cont = t_perm.contiguous()

print(f"squeeze: {t_sq.shape}")
print(f"permute: {t_perm.shape}")

a = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
b = torch.tensor([[5.0, 6.0], [7.0, 8.0]])

# 四则运算（逐元素）
print(a + b)     # 或 torch.add(a, b)
print(a - b)     # 或 torch.sub(a, b)
print(a * b)     # 逐元素乘（Hadamard product）
print(a / b)     # 逐元素除法
print(a ** 2)    # 逐元素幂运算

# 矩阵乘法
matmul = a @ b          # 或 torch.matmul(a, b)
mm = torch.mm(a, b)     # 仅限 2D

print(f"矩阵乘:\n{matmul}")

# 聚合运算
t = torch.rand(3, 4)
print(f"总和: {t.sum()}")
print(f"平均: {t.mean()}")
print(f"最大: {t.max()}")
print(f"最小: {t.min()}")
print(f"标准差: {t.std()}")

# 指定轴的聚合
print(f"按行求和: {t.sum(dim=0)}")  # 每一列的和（沿行方向）
print(f"按列求和: {t.sum(dim=1)}")  # 每一行的和（沿列方向）
print(f"keepdim:\n{t.sum(dim=1, keepdim=True)}")

# argmax/argmin
print(f"最大值索引: {t.argmax()}")
print(f"按行的最大值索引: {t.argmax(dim=1)}")

# 广播示例
a = torch.tensor([[1, 2, 3],
                  [4, 5, 6]])   # shape: [2, 3]
b = torch.tensor([10, 20, 30])  # shape: [3]

# b 自动扩展为 [2, 3] 后运算
print(a + b)
# tensor([[11, 22, 33],
#         [14, 25, 36]])

# 标量运算也支持广播
print(a * 2)   # 所有元素乘以 2
print(a + 100) # 所有元素加上 100

# 列向量 + 行向量
col = torch.tensor([[1], [2], [3]])  # shape: [3, 1]
row = torch.tensor([10, 20, 30])      # shape: [3]
print(col + row)  # shape: [3, 3] — 类似外积

t = torch.arange(24).reshape(2, 3, 4).float()

# 基本索引
print(t[0])        # 第一个矩阵（shape: [3, 4]）
print(t[0, 1])     # [3, 4] 矩阵的第二行（shape: [4]）
print(t[0, 1, 2])  # 标量

# 切片
print(t[:, 1:, :2])  # 全部、第 1 个之后、前 2 列

# 高级索引（Fancy indexing）
indices = torch.tensor([0, 2])
print(t[:, indices, :])  # 只选取第 0、2 行

# 条件索引（Boolean masking）
mask = t > 10
print(t[mask])  # 只提取大于 10 的元素（返回 1D 张量）

# where：根据条件从两个张量中选择
a = torch.tensor([1.0, 2.0, 3.0, 4.0])
b = torch.tensor([10.0, 20.0, 30.0, 40.0])
condition = a > 2
result = torch.where(condition, b, a)
print(result)  # tensor([ 1.,  2., 30., 40.])

import torch

# 创建 requires_grad=True 的张量 → 开始追踪运算
x = torch.tensor(3.0, requires_grad=True)
y = torch.tensor(4.0, requires_grad=True)

# 执行运算 → 构建计算图
z = x ** 2 + 2 * x * y + y ** 2  # (x + y)^2

print(f"z = {z}")  # z = 49.0

# 执行反向传播
z.backward()

# 查看梯度
# dz/dx = 2x + 2y = 2*3 + 2*4 = 14
print(f"dz/dx = {x.grad}")  # 14.0

# dz/dy = 2x + 2y = 14
print(f"dz/dy = {y.grad}")  # 14.0

# 向量函数的梯度
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = x ** 2  # [1, 4, 9]
z = y.sum()  # 归约为标量

z.backward()
print(f"x.grad: {x.grad}")  # [2, 4, 6]（dy/dx = 2x）

# gradient 参数：非标量的 backward
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = x ** 2  # [1, 4, 9]

# y.backward() 会报错 → y 不是标量
# 通过 gradient 参数传入加权和
grad_output = torch.tensor([1.0, 1.0, 1.0])  # 每个元素的权重
y.backward(gradient=grad_output)
print(f"x.grad: {x.grad}")  # [2, 4, 6]

# 梯度累积问题 — 需要初始化
x = torch.tensor(2.0, requires_grad=True)

for i in range(3):
    y = x ** 2
    y.backward()
    print(f"iteration {i}: x.grad = {x.grad}")
    # 每次不初始化的话就会累积
    x.grad.zero_()  # 原地初始化

# no_grad：推理时关闭梯度（节省内存）
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)

with torch.no_grad():
    y = x ** 2  # 不生成计算图
    print(f"y.requires_grad: {y.requires_grad}")  # False

# detach：从计算图中分离
x = torch.tensor([1.0, 2.0], requires_grad=True)
y = x * 2
z = y.detach()  # 分离梯度追踪
print(f"z.requires_grad: {z.requires_grad}")  # False

# 冻结部分参数（迁移学习时很有用）
for param in model.parameters():
    param.requires_grad = False

# 二阶微分示例
x = torch.tensor(3.0, requires_grad=True)
y = x ** 4

# 一阶微分：dy/dx = 4x^3
dy_dx = torch.autograd.grad(y, x, create_graph=True)[0]
print(f"一阶微分: {dy_dx}")  # 108

# 二阶微分：d2y/dx2 = 12x^2
d2y_dx2 = torch.autograd.grad(dy_dx, x)[0]
print(f"二阶微分: {d2y_dx2}")  # 108

import torch
import torch.nn as nn

# 定义一个简单模型
class SimpleModel(nn.Module):
    def __init__(self, in_features, hidden_size, out_features):
        super().__init__()
        # 定义层（参数自动注册）
        self.fc1 = nn.Linear(in_features, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, out_features)
        self.dropout = nn.Dropout(p=0.5)

    def forward(self, x):
        # 定义前向传播
        x = self.fc1(x)
        x = self.relu(x)
        x = self.dropout(x)
        x = self.fc2(x)
        return x

# 创建模型实例
model = SimpleModel(784, 256, 10)
print(model)

# 查看参数数量
total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"全部参数: {total_params:,}")
print(f"可训练参数: {trainable_params:,}")

# 访问参数
for name, param in model.named_parameters():
    print(f"{name}: {param.shape}")

# 执行前向传播
x = torch.randn(32, 784)  # batch_size=32, features=784
output = model(x)
print(f"输出 shape: {output.shape}")  # [32, 10]

# Sequential：按顺序构建层
seq_model = nn.Sequential(
    nn.Linear(784, 256),
    nn.ReLU(),
    nn.Dropout(0.3),
    nn.Linear(256, 128),
    nn.ReLU(),
    nn.Linear(128, 10)
)

# ModuleList：用列表管理层
class ResidualBlock(nn.Module):
    def __init__(self, num_blocks, hidden_size):
        super().__init__()
        self.layers = nn.ModuleList([
            nn.Linear(hidden_size, hidden_size)
            for _ in range(num_blocks)
        ])
        self.relu = nn.ReLU()

    def forward(self, x):
        for layer in self.layers:
            x = self.relu(layer(x)) + x  # 残差连接
        return x

# ModuleDict：用字典管理层
class MultiTaskModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.backbone = nn.Linear(784, 256)
        self.heads = nn.ModuleDict({
            'classification': nn.Linear(256, 10),
            'regression': nn.Linear(256, 1)
        })

    def forward(self, x, task='classification'):
        features = torch.relu(self.backbone(x))
        return self.heads[task](features)

import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt

# 生成数据
torch.manual_seed(42)
n_samples = 200

# y = 3x + 2 + 噪声
X = torch.linspace(-5, 5, n_samples).unsqueeze(1)  # [200, 1]
y_true = 3 * X + 2
y = y_true + torch.randn_like(y_true) * 0.5       # 添加噪声

# 定义模型
class LinearRegression(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(1, 1)

    def forward(self, x):
        return self.linear(x)

model = LinearRegression()

# 损失函数与优化器
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 训练循环
n_epochs = 1000
losses = []

for epoch in range(n_epochs):
    # 1. 前向传播（Forward Pass）
    y_pred = model(X)

    # 2. 计算损失
    loss = criterion(y_pred, y)
    losses.append(loss.item())

    # 3. 梯度归零（重要！）
    optimizer.zero_grad()

    # 4. 反向传播（Backward Pass）
    loss.backward()

    # 5. 更新参数
    optimizer.step()

    if (epoch + 1) % 200 == 0:
        w = model.linear.weight.item()
        b = model.linear.bias.item()
        print(f"Epoch {epoch+1}: Loss={loss.item():.4f}, w={w:.4f}, b={b:.4f}")

# 查看结果
print(f"\n训练得到的权重: {model.linear.weight.item():.4f}（正确答案：3.0）")
print(f"训练得到的偏置: {model.linear.bias.item():.4f}（正确答案：2.0）")

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# 超参数
BATCH_SIZE = 64
LEARNING_RATE = 0.001
N_EPOCHS = 10
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 数据预处理与加载
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))  # MNIST 的均值、标准差
])

train_dataset = datasets.MNIST(root='./data', train=True,
                                download=True, transform=transform)
test_dataset = datasets.MNIST(root='./data', train=False,
                               transform=transform)

train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE,
                          shuffle=True, num_workers=2)
test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE,
                         shuffle=False, num_workers=2)

# 定义 MLP 模型
class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.network = nn.Sequential(
            nn.Flatten(),              # 28x28 → 784
            nn.Linear(784, 512),
            nn.BatchNorm1d(512),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(512, 256),
            nn.BatchNorm1d(256),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(256, 10)
        )

    def forward(self, x):
        return self.network(x)

model = MLP().to(DEVICE)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)

# 训练函数
def train_epoch(model, loader, criterion, optimizer, device):
    model.train()
    total_loss = 0
    correct = 0
    total = 0

    for batch_idx, (data, target) in enumerate(loader):
        data, target = data.to(device), target.to(device)

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

        total_loss += loss.item()
        pred = output.argmax(dim=1)
        correct += pred.eq(target).sum().item()
        total += target.size(0)

    return total_loss / len(loader), 100.0 * correct / total

# 评估函数
def evaluate(model, loader, criterion, device):
    model.eval()
    total_loss = 0
    correct = 0
    total = 0

    with torch.no_grad():
        for data, target in loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            total_loss += criterion(output, target).item()
            pred = output.argmax(dim=1)
            correct += pred.eq(target).sum().item()
            total += target.size(0)

    return total_loss / len(loader), 100.0 * correct / total

# 执行训练
for epoch in range(N_EPOCHS):
    train_loss, train_acc = train_epoch(model, train_loader, criterion, optimizer, DEVICE)
    test_loss, test_acc = evaluate(model, test_loader, criterion, DEVICE)
    print(f"Epoch {epoch+1}/{N_EPOCHS} | "
          f"Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.2f}% | "
          f"Test Loss: {test_loss:.4f}, Test Acc: {test_acc:.2f}%")

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# 准备 CIFAR-10 数据
transform_train = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(brightness=0.2, contrast=0.2),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465),
                         (0.2023, 0.1994, 0.2010))
])

transform_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465),
                         (0.2023, 0.1994, 0.2010))
])

train_data = datasets.CIFAR10('./data', train=True, download=True, transform=transform_train)
test_data = datasets.CIFAR10('./data', train=False, transform=transform_test)

train_loader = DataLoader(train_data, batch_size=128, shuffle=True, num_workers=4)
test_loader = DataLoader(test_data, batch_size=128, shuffle=False, num_workers=4)

CLASSES = ['airplane', 'automobile', 'bird', 'cat', 'deer',
           'dog', 'frog', 'horse', 'ship', 'truck']

# 定义 CNN 模型（VGG 风格）
class CNN(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        # 特征提取部分
        self.features = nn.Sequential(
            # Block 1: 3 → 64
            nn.Conv2d(3, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2),      # 32x32 → 16x16
            nn.Dropout2d(0.1),

            # Block 2: 64 → 128
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2),      # 16x16 → 8x8
            nn.Dropout2d(0.2),

            # Block 3: 128 → 256
            nn.Conv2d(128, 256, kernel_size=3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2),      # 8x8 → 4x4
        )

        # 分类部分
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(256 * 4 * 4, 1024),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Linear(1024, 512),
            nn.ReLU(inplace=True),
            nn.Dropout(0.3),
            nn.Linear(512, num_classes)
        )

    def forward(self, x):
        x = self.features(x)
        x = self.classifier(x)
        return x

model = CNN().to(DEVICE)
print(f"模型参数量: {sum(p.numel() for p in model.parameters()):,}")

import torch
import torch.nn as nn
import numpy as np

# 基于 LSTM 的时间序列预测模型
class LSTMPredictor(nn.Module):
    def __init__(self, input_size=1, hidden_size=64, num_layers=2,
                 output_size=1, dropout=0.2):
        super().__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers

        # LSTM 层
        self.lstm = nn.LSTM(
            input_size=input_size,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True,     # 输入: [batch, seq_len, features]
            dropout=dropout if num_layers > 1 else 0,
            bidirectional=False
        )

        # 输出层
        self.fc = nn.Sequential(
            nn.Linear(hidden_size, 32),
            nn.ReLU(),
            nn.Linear(32, output_size)
        )

    def forward(self, x):
        # x shape: [batch_size, seq_len, input_size]
        batch_size = x.size(0)

        # 初始 hidden/cell state（初始化为 0）
        h0 = torch.zeros(self.num_layers, batch_size,
                         self.hidden_size).to(x.device)
        c0 = torch.zeros(self.num_layers, batch_size,
                         self.hidden_size).to(x.device)

        # LSTM 前向传播
        # out: [batch_size, seq_len, hidden_size]
        out, (hn, cn) = self.lstm(x, (h0, c0))

        # 只使用最后一个时间步的输出
        out = self.fc(out[:, -1, :])  # [batch_size, output_size]
        return out

# 以正弦波数据为例
t = np.linspace(0, 100, 1000)
data = np.sin(0.5 * t) + 0.1 * np.random.randn(1000)
data = torch.FloatTensor(data).unsqueeze(1)

# 生成序列数据的函数
def create_sequences(data, seq_len=50):
    X, y = [], []
    for i in range(len(data) - seq_len):
        X.append(data[i:i+seq_len])
        y.append(data[i+seq_len])
    return torch.stack(X), torch.stack(y)

X, y = create_sequences(data, seq_len=50)
print(f"X shape: {X.shape}")  # [950, 50, 1]
print(f"y shape: {y.shape}")  # [950, 1]

# GRU — 参数比 LSTM 更少的变体
class GRUPredictor(nn.Module):
    def __init__(self, input_size=1, hidden_size=64, num_layers=2):
        super().__init__()
        self.gru = nn.GRU(input_size, hidden_size, num_layers,
                          batch_first=True, dropout=0.2)
        self.fc = nn.Linear(hidden_size, 1)

    def forward(self, x):
        out, _ = self.gru(x)
        return self.fc(out[:, -1, :])

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

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model=512, num_heads=8, dropout=0.1):
        super().__init__()
        assert d_model % num_heads == 0

        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads  # 每个头的维度

        # Q、K、V、输出的投影
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)

        self.dropout = nn.Dropout(dropout)
        self.scale = math.sqrt(self.d_k)

    def split_heads(self, x):
        # x: [batch, seq, d_model] → [batch, num_heads, seq, d_k]
        batch, seq, _ = x.shape
        x = x.view(batch, seq, self.num_heads, self.d_k)
        return x.transpose(1, 2)

    def forward(self, query, key, value, mask=None):
        batch_size = query.size(0)

        # 1. Q、K、V 线性变换并拆分头
        Q = self.split_heads(self.W_q(query))  # [B, H, Sq, dk]
        K = self.split_heads(self.W_k(key))    # [B, H, Sk, dk]
        V = self.split_heads(self.W_v(value))  # [B, H, Sk, dk]

        # 2. Scaled Dot-Product Attention
        scores = torch.matmul(Q, K.transpose(-2, -1)) / self.scale
        # scores: [B, H, Sq, Sk]

        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))

        attn_weights = F.softmax(scores, dim=-1)
        attn_weights = self.dropout(attn_weights)

        # 3. 与 Value 的加权和
        context = torch.matmul(attn_weights, V)  # [B, H, Sq, dk]

        # 4. 合并多头
        context = context.transpose(1, 2).contiguous()
        context = context.view(batch_size, -1, self.d_model)

        # 5. 输出投影
        output = self.W_o(context)
        return output, attn_weights

class FeedForward(nn.Module):
    def __init__(self, d_model=512, d_ff=2048, dropout=0.1):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model)
        )

    def forward(self, x):
        return self.net(x)

class TransformerEncoderLayer(nn.Module):
    def __init__(self, d_model=512, num_heads=8, d_ff=2048, dropout=0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads, dropout)
        self.ff = FeedForward(d_model, d_ff, dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        # Pre-Norm 方式（原论文是 Post-Norm）
        attn_out, _ = self.self_attn(x, x, x, mask)
        x = self.norm1(x + self.dropout(attn_out))

        ff_out = self.ff(x)
        x = self.norm2(x + self.dropout(ff_out))
        return x

class PositionalEncoding(nn.Module):
    def __init__(self, d_model=512, max_len=5000, dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout)

        # 计算位置编码
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1).float()
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
        )
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)  # [1, max_len, d_model]
        self.register_buffer('pe', pe)

    def forward(self, x):
        x = x + self.pe[:, :x.size(1)]
        return self.dropout(x)

# 使用示例
d_model = 512
encoder_layer = TransformerEncoderLayer(d_model=d_model, num_heads=8)
pos_enc = PositionalEncoding(d_model=d_model)

x = torch.randn(2, 10, d_model)  # [batch=2, seq=10, d_model=512]
x = pos_enc(x)
output = encoder_layer(x)
print(f"Transformer Encoder 输出: {output.shape}")  # [2, 10, 512]

import torch
from torch.utils.data import Dataset, DataLoader
import pandas as pd
import numpy as np
from PIL import Image
import os

# 自定义 Dataset 实现
class CustomImageDataset(Dataset):
    def __init__(self, csv_file, img_dir, transform=None):
        """
        csv_file: 包含图像路径和标签的 CSV
        img_dir: 图像根目录
        transform: torchvision transforms
        """
        self.annotations = pd.read_csv(csv_file)
        self.img_dir = img_dir
        self.transform = transform

    def __len__(self):
        # 返回数据集大小（必需）
        return len(self.annotations)

    def __getitem__(self, idx):
        # 根据索引返回样本（必需）
        img_path = os.path.join(self.img_dir, self.annotations.iloc[idx, 0])
        image = Image.open(img_path).convert('RGB')
        label = int(self.annotations.iloc[idx, 1])

        if self.transform:
            image = self.transform(image)

        return image, label

# 用于数值数据的 Dataset
class TabularDataset(Dataset):
    def __init__(self, X, y):
        self.X = torch.FloatTensor(X)
        self.y = torch.LongTensor(y)

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

    def __getitem__(self, idx):
        return self.X[idx], self.y[idx]

# DataLoader 高级用法
dataset = TabularDataset(
    X=np.random.randn(1000, 20),
    y=np.random.randint(0, 5, 1000)
)

# 基础 DataLoader
basic_loader = DataLoader(dataset, batch_size=32, shuffle=True)

# 高级设置
advanced_loader = DataLoader(
    dataset,
    batch_size=64,
    shuffle=True,
    num_workers=4,        # 并行数据加载（按 CPU 核心数调整）
    pin_memory=True,      # 加速 GPU 传输（使用 CUDA 时）
    drop_last=True,       # 丢弃最后一个不完整的 batch
    prefetch_factor=2,    # 预取的 batch 数
    persistent_workers=True  # 复用 worker 进程
)

# 查看 batch
for batch_X, batch_y in advanced_loader:
    print(f"batch X: {batch_X.shape}")  # [64, 20]
    print(f"batch y: {batch_y.shape}")  # [64]
    break

# WeightedRandomSampler：处理类别不平衡
from torch.utils.data import WeightedRandomSampler

class_counts = [800, 150, 50]  # 各类别的样本数
weights = 1.0 / torch.tensor(class_counts, dtype=torch.float)
# 为每个样本分配类别权重
sample_weights = weights[dataset.y]  # 每个样本的权重

sampler = WeightedRandomSampler(
    weights=sample_weights,
    num_samples=len(dataset),
    replacement=True
)

balanced_loader = DataLoader(dataset, batch_size=32, sampler=sampler)

import torch.optim as optim

# 示例模型
model = nn.Linear(100, 10)

# SGD（Stochastic Gradient Descent）
sgd = optim.SGD(
    model.parameters(),
    lr=0.01,
    momentum=0.9,      # 保持上一次更新的方向
    weight_decay=1e-4, # L2 正则化
    nesterov=True      # Nesterov momentum
)

# Adam：自适应学习率
adam = optim.Adam(
    model.parameters(),
    lr=0.001,
    betas=(0.9, 0.999),  # 一阶、二阶矩的衰减率
    eps=1e-8,
    weight_decay=0
)

# AdamW：Adam + 正确的 Weight Decay
# 注意：Adam 的 weight_decay 与 L2 正则化不同
# Transformer 系列模型推荐使用 AdamW
adamw = optim.AdamW(
    model.parameters(),
    lr=1e-3,
    betas=(0.9, 0.999),
    weight_decay=0.01   # 与学习率无关地独立施加
)

# RMSprop：对循环神经网络效果良好
rmsprop = optim.RMSprop(
    model.parameters(),
    lr=0.01,
    alpha=0.99,
    momentum=0.0
)

# 按参数组设置不同学习率（迁移学习时很有用）
optimizer = optim.Adam([
    {'params': model.features.parameters(), 'lr': 1e-4},  # 主干网络：较低 LR
    {'params': model.classifier.parameters(), 'lr': 1e-3} # 分类头：较高 LR
], lr=1e-3)

# 保存/恢复优化器状态
checkpoint = {
    'model': model.state_dict(),
    'optimizer': optimizer.state_dict(),
    'epoch': 10
}
torch.save(checkpoint, 'checkpoint.pt')

# 恢复
ckpt = torch.load('checkpoint.pt')
model.load_state_dict(ckpt['model'])
optimizer.load_state_dict(ckpt['optimizer'])
start_epoch = ckpt['epoch']

import torch.optim as optim
from torch.optim.lr_scheduler import (
    StepLR, MultiStepLR, ExponentialLR,
    CosineAnnealingLR, OneCycleLR,
    ReduceLROnPlateau, CosineAnnealingWarmRestarts
)

model = nn.Linear(10, 2)
optimizer = optim.SGD(model.parameters(), lr=0.1)

# StepLR：每 step_size 个 epoch 就衰减 gamma 倍
step_scheduler = StepLR(optimizer, step_size=30, gamma=0.1)
# 0~29: lr=0.1, 30~59: lr=0.01, 60~89: lr=0.001

# MultiStepLR：在指定 epoch 衰减
multi_scheduler = MultiStepLR(optimizer, milestones=[50, 100, 150], gamma=0.1)

# CosineAnnealingLR：按余弦周期衰减
cosine_scheduler = CosineAnnealingLR(optimizer, T_max=100, eta_min=1e-6)

# ReduceLROnPlateau：验证损失不再改善时衰减（最实用）
plateau_scheduler = ReduceLROnPlateau(
    optimizer,
    mode='min',       # 以最小化为目标（loss）
    factor=0.5,       # 衰减比例
    patience=10,      # 允许多少个 epoch 无改善
    min_lr=1e-7,
    verbose=True
)

# OneCycleLR：快速收敛（super convergence）
one_cycle = OneCycleLR(
    optimizer,
    max_lr=0.01,
    steps_per_epoch=100,  # len(train_loader)
    epochs=30,
    pct_start=0.3,        # 用整体的 30% 做 warm-up
    anneal_strategy='cos'
)

# CosineAnnealingWarmRestarts：通过 warm restart 周期性重置
warm_restart = CosineAnnealingWarmRestarts(
    optimizer,
    T_0=10,    # 到第一次重置为止的 epoch 数
    T_mult=2,  # 每次重置后周期变为 T_mult 倍
    eta_min=1e-6
)

# 在训练循环中使用调度器
for epoch in range(100):
    train_loss = 0.5  # 实际训练循环的结果

    # 大多数调度器：按 epoch 调用 step
    cosine_scheduler.step()

    # ReduceLROnPlateau：把验证指标作为参数传入
    plateau_scheduler.step(train_loss)

    # OneCycleLR：按 batch 调用 step
    # for batch in loader:
    #     ...
    #     one_cycle.step()

    print(f"Epoch {epoch+1}: LR = {optimizer.param_groups[0]['lr']:.6f}")

import torch.nn as nn

# Dropout：训练时随机让神经元失效
class DropoutDemo(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(100, 50)
        self.dropout = nn.Dropout(p=0.5)  # 50% 失效
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.dropout(x)  # 仅训练时生效，评估时自动关闭
        return self.fc2(x)

# model.train() → 激活 Dropout
# model.eval() → 关闭 Dropout

# BatchNorm1d：对 mini-batch 做归一化（用于 FC 层之后）
bn_model = nn.Sequential(
    nn.Linear(100, 64),
    nn.BatchNorm1d(64),  # 沿 batch 维度归一化
    nn.ReLU(),
    nn.Linear(64, 32),
    nn.BatchNorm1d(32),
    nn.ReLU(),
    nn.Linear(32, 10)
)

# BatchNorm2d：2D feature map（用于 CNN 层之后）
cnn_with_bn = nn.Sequential(
    nn.Conv2d(3, 32, 3, padding=1),
    nn.BatchNorm2d(32),
    nn.ReLU(),
    nn.Conv2d(32, 64, 3, padding=1),
    nn.BatchNorm2d(64),
    nn.ReLU()
)

# LayerNorm：沿特征维度归一化（主要用于 Transformer）
# 与 BatchNorm 不同，不依赖 batch 大小
transformer_norm = nn.Sequential(
    nn.Linear(512, 512),
    nn.LayerNorm(512),  # 沿最后一维（512）归一化
    nn.ReLU()
)

# GroupNorm：BatchNorm 和 LayerNorm 之间的折中（对小 batch 有用）
group_norm = nn.GroupNorm(
    num_groups=8,    # 把通道分成 8 组
    num_channels=64  # 总通道数
)

# InstanceNorm：常用于风格迁移等场景
instance_norm = nn.InstanceNorm2d(64)

# 各归一化方法对比小结：
# BatchNorm  : 沿 batch × 空间归一化 → 对 CNN 有效，依赖 batch 大小
# LayerNorm  : 沿特征维度归一化 → 对 Transformer、RNN 有效
# GroupNorm  : 小 batch 场景下 BatchNorm 的替代方案
# InstanceNorm: 用于风格迁移、图像生成

import torchvision.models as models
import torch.nn as nn

# 加载预训练模型
# 推荐通过 weights 参数显式指定（最新 API）
resnet50 = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V2)
vgg16 = models.vgg16(weights=models.VGG16_Weights.IMAGENET1K_V1)
efficientnet = models.efficientnet_b0(weights=models.EfficientNet_B0_Weights.IMAGENET1K_V1)
vit = models.vit_b_16(weights=models.ViT_B_16_Weights.IMAGENET1K_V1)

print(resnet50)  # 输出结构

# 方法 1：Feature Extractor（冻结主干网络）
# 固定主干网络参数
for param in resnet50.parameters():
    param.requires_grad = False

# 只替换最后的 FC 层（匹配新的类别数）
num_classes = 5
resnet50.fc = nn.Linear(resnet50.fc.in_features, num_classes)

# 只有最后一层会被训练
trainable = sum(p.numel() for p in resnet50.parameters() if p.requires_grad)
print(f"可训练参数: {trainable:,}")  # 约 2,050 个

# 方法 2：Fine-tuning（训练全部或部分层）
resnet_ft = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V2)
resnet_ft.fc = nn.Linear(resnet_ft.fc.in_features, num_classes)

# 按层设置不同学习率（越靠底层 = 越低的 LR）
optimizer = torch.optim.AdamW([
    {'params': resnet_ft.layer1.parameters(), 'lr': 1e-5},
    {'params': resnet_ft.layer2.parameters(), 'lr': 1e-5},
    {'params': resnet_ft.layer3.parameters(), 'lr': 1e-4},
    {'params': resnet_ft.layer4.parameters(), 'lr': 1e-4},
    {'params': resnet_ft.fc.parameters(),     'lr': 1e-3},
], lr=1e-4, weight_decay=0.01)

# 方法 3：用 torchvision transforms 做数据预处理
from torchvision import transforms

# 预训练模型的输入归一化数值（基于 ImageNet）
preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

import torch
import torch.nn as nn

model = nn.Linear(10, 5)
optimizer = torch.optim.Adam(model.parameters())

# 方法 1：保存 state_dict（推荐）
# 只保存模型参数（不含结构）
torch.save(model.state_dict(), 'model_weights.pt')

# 加载
loaded_model = nn.Linear(10, 5)  # 需要相同的架构
loaded_model.load_state_dict(torch.load('model_weights.pt',
                                         weights_only=True))
loaded_model.eval()

# 方法 2：保存整个模型（不推荐——可移植性较差）
torch.save(model, 'full_model.pt')
loaded_full = torch.load('full_model.pt', weights_only=False)

# 方法 3：checkpoint — 为恢复训练保存完整状态
def save_checkpoint(model, optimizer, scheduler, epoch, loss, path):
    torch.save({
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'scheduler_state_dict': scheduler.state_dict() if scheduler else None,
        'loss': loss,
    }, path)
    print(f"已保存 checkpoint: {path}")

def load_checkpoint(path, model, optimizer=None, scheduler=None):
    checkpoint = torch.load(path, map_location='cpu', weights_only=True)
    model.load_state_dict(checkpoint['model_state_dict'])

    if optimizer:
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    if scheduler and checkpoint['scheduler_state_dict']:
        scheduler.load_state_dict(checkpoint['scheduler_state_dict'])

    return checkpoint['epoch'], checkpoint['loss']

# 使用示例
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)
save_checkpoint(model, optimizer, scheduler, epoch=50, loss=0.25,
                path='checkpoint_ep50.pt')

start_epoch, prev_loss = load_checkpoint('checkpoint_ep50.pt',
                                          model, optimizer, scheduler)
print(f"恢复: epoch={start_epoch}, loss={prev_loss:.4f}")

# 把 GPU 模型加载到 CPU
model_cpu = nn.Linear(10, 5)
model_cpu.load_state_dict(
    torch.load('model_weights.pt', map_location='cpu', weights_only=True)
)

import torch
import torch.nn as nn

class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(10, 5)

    def forward(self, x):
        return torch.relu(self.fc(x))

model = SimpleNet()
model.eval()

# 方法 1：torch.jit.script — 编译整个模型
scripted_model = torch.jit.script(model)

# 保存与加载
scripted_model.save('model_scripted.pt')
loaded_scripted = torch.jit.load('model_scripted.pt')

x = torch.randn(4, 10)
with torch.no_grad():
    out = loaded_scripted(x)
print(f"TorchScript 输出: {out.shape}")

# 方法 2：torch.jit.trace — 用示例输入进行追踪
example_input = torch.randn(1, 10)
traced_model = torch.jit.trace(model, example_input)
traced_model.save('model_traced.pt')

# 方法 3：导出 ONNX（兼容其他框架）
import torch.onnx

dummy_input = torch.randn(1, 10)
torch.onnx.export(
    model,
    dummy_input,
    'model.onnx',
    export_params=True,
    opset_version=17,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)
print("ONNX 导出完成")

# torch.compile（PyTorch 2.0+）：最新的编译方式
# 无需改动既有代码即可应用
compiled_model = torch.compile(model)
out = compiled_model(x)
print(f"torch.compile 输出: {out.shape}")

# train_ddp.py — 编写为独立执行脚本
import os
import torch
import torch.nn as nn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler
from torchvision import datasets, transforms

def setup(rank, world_size):
    """初始化进程组"""
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'

    # nccl: GPU 通信后端（推荐）
    # gloo: 用于 CPU 或调试
    dist.init_process_group(
        backend='nccl',
        rank=rank,
        world_size=world_size
    )

def cleanup():
    """清理进程组"""
    dist.destroy_process_group()

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(784, 256),
            nn.ReLU(),
            nn.Linear(256, 10)
        )

    def forward(self, x):
        return self.net(x.view(x.size(0), -1))

def train(rank, world_size, num_epochs=5):
    print(f"进程 {rank}/{world_size} 启动")
    setup(rank, world_size)

    # 为每个进程分配 GPU
    torch.cuda.set_device(rank)
    device = torch.device(f'cuda:{rank}')

    # 创建模型并用 DDP 包装
    model = SimpleModel().to(device)
    ddp_model = DDP(model, device_ids=[rank])

    # 数据集与分布式采样器
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    dataset = datasets.MNIST('./data', train=True, download=True,
                              transform=transform)

    # DistributedSampler：给每个进程分配不同的数据
    sampler = DistributedSampler(
        dataset,
        num_replicas=world_size,
        rank=rank,
        shuffle=True
    )

    loader = DataLoader(
        dataset,
        batch_size=128,
        sampler=sampler,
        num_workers=4,
        pin_memory=True
    )

    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(ddp_model.parameters(), lr=0.001)

    for epoch in range(num_epochs):
        # 每个 epoch 都更新采样器的种子（打乱数据）
        sampler.set_epoch(epoch)

        ddp_model.train()
        total_loss = 0.0

        for batch_idx, (data, target) in enumerate(loader):
            data, target = data.to(device), target.to(device)

            optimizer.zero_grad()
            output = ddp_model(data)
            loss = criterion(output, target)
            loss.backward()  # 自动同步梯度
            optimizer.step()

            total_loss += loss.item()

        # 只在 rank 0 上输出日志
        if rank == 0:
            avg_loss = total_loss / len(loader)
            print(f"Epoch {epoch+1}: Average Loss = {avg_loss:.4f}")

    cleanup()

# 执行方式: torchrun --nproc_per_node=4 train_ddp.py
if __name__ == '__main__':
    import torch.multiprocessing as mp
    world_size = torch.cuda.device_count()
    mp.spawn(
        train,
        args=(world_size, 5),
        nprocs=world_size,
        join=True
    )

# 单节点 4 块 GPU 训练
torchrun --nproc_per_node=4 train_ddp.py

# 多节点训练（节点 0）
torchrun --nnodes=2 --nproc_per_node=4 \
         --node_rank=0 \
         --master_addr="192.168.1.100" \
         --master_port=12355 \
         train_ddp.py

# DataParallel (DP)：简单但效率不高
# - 所有梯度都汇集到 GPU 0 → 瓶颈
# - 采用多线程方式而非多进程
model_dp = nn.DataParallel(model, device_ids=[0, 1, 2, 3])

# DistributedDataParallel (DDP)：推荐方式
# - 每块 GPU 独立计算梯度
# - 通过 All-Reduce 高效同步
# - 即使单 GPU 上 DDP 也更快（规避 Python GIL）
model_ddp = DDP(model, device_ids=[rank])

from torch.cuda.amp import autocast, GradScaler

model = SimpleModel().to('cuda')
optimizer = torch.optim.Adam(model.parameters())
scaler = GradScaler()  # FP16 损失缩放

for data, target in train_loader:
    data, target = data.to('cuda'), target.to('cuda')
    optimizer.zero_grad()

    # 用 FP16 做前向传播
    with autocast():
        output = model(data)
        loss = criterion(output, target)

    # 缩放后的反向传播
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

# 防止梯度爆炸
max_grad_norm = 1.0
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
optimizer.step()

import random
import numpy as np

def set_seed(seed=42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    # 为了完全可复现（会有性能损耗）
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

set_seed(42)