💡 왼쪽 원문을 읽으면서 오른쪽에 따라 써보세요. Tab 키로 힌트를 받을 수 있습니다.

PyTorch 高级技巧完全指南

PyTorch 高级技巧完全指南

PyTorch 是研究与生产部署领域都最受欢迎的深度学习框架之一。但大多数开发者止步于基础的张量运算和 nn.Module 定义。本指南将介绍在实际生产环境中最大化性能、实现自定义运算、以及高效管理内存的高级技巧。

1. torch.compile（PyTorch 2.0+）

PyTorch 2.0 引入的 torch.compile 通过编译模型来大幅提升执行速度。与传统的 TorchScript 或 ONNX export 不同，torch.compile 几乎不需要修改 Python 代码，就能带来 2 倍以上的速度提升。

torch.compile 简介与优势

torch.compile 由三个核心组件构成：

TorchDynamo：拦截 Python 字节码来生成 FX 图
AOTAutograd：预先编译自动微分图
Inductor：通过 TorchInductor 后端生成优化后的内核（Triton GPU 内核或 C++ CPU 内核）

import torch
import torch.nn as nn
import time

# 定义基础模型
class TransformerBlock(nn.Module):
    def __init__(self, d_model=512, nhead=8, dim_feedforward=2048):
        super().__init__()
        self.attention = nn.MultiheadAttention(d_model, nhead, batch_first=True)
        self.feed_forward = nn.Sequential(
            nn.Linear(d_model, dim_feedforward),
            nn.ReLU(),
            nn.Linear(dim_feedforward, d_model)
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)

    def forward(self, x):
        attn_out, _ = self.attention(x, x, x)
        x = self.norm1(x + attn_out)
        ff_out = self.feed_forward(x)
        x = self.norm2(x + ff_out)
        return x

device = "cuda" if torch.cuda.is_available() else "cpu"
model = TransformerBlock().to(device)

# 应用 torch.compile
compiled_model = torch.compile(model)

# 预热
x = torch.randn(32, 128, 512, device=device)
for _ in range(3):
    _ = compiled_model(x)

# 性能比较
N = 100

# 常规模型
start = time.perf_counter()
for _ in range(N):
    _ = model(x)
if device == "cuda":
    torch.cuda.synchronize()
elapsed_eager = time.perf_counter() - start

# 编译后的模型
start = time.perf_counter()
for _ in range(N):
    _ = compiled_model(x)
if device == "cuda":
    torch.cuda.synchronize()
elapsed_compiled = time.perf_counter() - start

print(f"Eager mode: {elapsed_eager:.3f}s")
print(f"Compiled mode: {elapsed_compiled:.3f}s")
print(f"Speedup: {elapsed_eager / elapsed_compiled:.2f}x")

编译模式（Compilation Modes）

torch.compile 提供三种模式：

# 默认模式 - 编译快，性能好
model_default = torch.compile(model, mode="default")

# 减少开销模式 - 更适合小模型
model_reduce = torch.compile(model, mode="reduce-overhead")

# 最大自动调优 - 编译时间长，性能最佳
model_autotune = torch.compile(model, mode="max-autotune")

# 全图模式 - 不允许动态图（strict）
model_full = torch.compile(model, fullgraph=True)

# 后端选择
model_eager = torch.compile(model, backend="eager")    # 不编译（用于调试）
model_aot = torch.compile(model, backend="aot_eager")  # 仅应用 AOT
model_inductor = torch.compile(model, backend="inductor")  # 默认值

动态形状支持

import torch._dynamo as dynamo

# 启用动态形状
model_dynamic = torch.compile(model, dynamic=True)

# 不同批大小下也无需重新编译
for batch_size in [8, 16, 32, 64]:
    x = torch.randn(batch_size, 128, 512, device=device)
    out = model_dynamic(x)
    print(f"Batch {batch_size}: output shape {out.shape}")

# 检查编译缓存
print(dynamo.explain(model)(x))

迁移已有代码

# 在训练循环中应用 torch.compile
def train_epoch(model, optimizer, dataloader, criterion):
    model.train()
    total_loss = 0

    for batch_idx, (data, target) in enumerate(dataloader):
        data, target = data.cuda(), target.cuda()
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()

    return total_loss / len(dataloader)

# 只需编译模型即可
model = MyModel().cuda()
compiled_model = torch.compile(model)  # 只需加这一行！

optimizer = torch.optim.Adam(compiled_model.parameters())
criterion = nn.CrossEntropyLoss()

2. 自定义自动微分（Custom Autograd）

PyTorch 的自动微分引擎十分强大，但有时仍需要数值上更稳定或更高效的自定义梯度计算。

继承 torch.autograd.Function

import torch
from torch.autograd import Function

class SigmoidFunction(Function):
    """数值稳定的 Sigmoid 实现"""

    @staticmethod
    def forward(ctx, input):
        # sigmoid = 1 / (1 + exp(-x))
        sigmoid = torch.sigmoid(input)
        # 保存 backward 中要用到的张量
        ctx.save_for_backward(sigmoid)
        return sigmoid

    @staticmethod
    def backward(ctx, grad_output):
        # 取出保存的 sigmoid
        sigmoid, = ctx.saved_tensors
        # gradient = sigmoid * (1 - sigmoid) * grad_output
        grad_input = sigmoid * (1 - sigmoid) * grad_output
        return grad_input

# 使用示例
def custom_sigmoid(x):
    return SigmoidFunction.apply(x)

x = torch.randn(3, 4, requires_grad=True)
y = custom_sigmoid(x)
loss = y.sum()
loss.backward()
print(f"Gradient: {x.grad}")

更复杂的示例：Leaky ReLU 与自定义 Backward

class LeakyReLUFunction(Function):
    @staticmethod
    def forward(ctx, input, negative_slope=0.01):
        ctx.save_for_backward(input)
        ctx.negative_slope = negative_slope
        return input.clamp(min=0) + negative_slope * input.clamp(max=0)

    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        negative_slope = ctx.negative_slope
        grad_input = grad_output.clone()
        grad_input[input < 0] *= negative_slope
        # 第二个输入（negative_slope）对应的梯度为 None
        return grad_input, None

class CustomLeakyReLU(nn.Module):
    def __init__(self, negative_slope=0.01):
        super().__init__()
        self.negative_slope = negative_slope

    def forward(self, x):
        return LeakyReLUFunction.apply(x, self.negative_slope)

数值梯度检查

from torch.autograd import gradcheck

def test_custom_op():
    # 数值梯度检查（推荐使用 float64）
    input = torch.randn(3, 4, dtype=torch.float64, requires_grad=True)

    # gradcheck 会数值计算雅可比矩阵，并与自动微分结果比较
    result = gradcheck(SigmoidFunction.apply, (input,), eps=1e-6, atol=1e-4)
    print(f"Gradient check passed: {result}")

    # 二阶反向传播测试
    input = torch.randn(3, 4, dtype=torch.float64, requires_grad=True)
    result = gradcheck(
        SigmoidFunction.apply,
        (input,),
        eps=1e-6,
        atol=1e-4,
        check_grad_dtypes=True
    )

test_custom_op()

二阶反向传播（Double Backpropagation）

class SquaredFunction(Function):
    """x^2 的自定义实现 - 支持二阶反向传播"""

    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)
        return x ** 2

    @staticmethod
    def backward(ctx, grad_output):
        x, = ctx.saved_tensors
        # 要支持二阶反向传播，需要设置 create_graph=True
        return 2 * x * grad_output

# 二阶反向传播示例（可用于 MAML 等元学习）
x = torch.randn(3, requires_grad=True)
y = SquaredFunction.apply(x)
grad_x = torch.autograd.grad(y.sum(), x, create_graph=True)[0]
# grad_x = 2x，对其再次反向传播
grad_grad_x = torch.autograd.grad(grad_x.sum(), x)[0]
# grad_grad_x = 2（常数）
print(f"Second derivative: {grad_grad_x}")

3. 自定义 CUDA 算子

torch.utils.cpp_extension 简介

PyTorch 提供了编写 C++/CUDA 扩展、并在 Python 中调用它们的工具。

# JIT 编译方式（适合开发/原型阶段）
from torch.utils.cpp_extension import load_inline

# C++ CPU 算子
cpp_source = """
#include <torch/extension.h>

torch::Tensor relu_forward(torch::Tensor input) {
    return input.clamp_min(0);
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("relu_forward", &relu_forward, "ReLU forward");
}
"""

# 内联编译
custom_relu_cpp = load_inline(
    name="custom_relu_cpp",
    cpp_sources=cpp_source,
    functions=["relu_forward"],
    verbose=False
)

x = torch.randn(5)
result = custom_relu_cpp.relu_forward(x)
print(result)

CUDA 内核示例：Fused Softmax

# CUDA 源码
cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

__global__ void fused_softmax_kernel(
    const float* input,
    float* output,
    int rows,
    int cols
) {
    int row = blockIdx.x;
    if (row >= rows) return;

    const float* row_input = input + row * cols;
    float* row_output = output + row * cols;

    // 求最大值（用于数值稳定性）
    float max_val = row_input[0];
    for (int i = 1; i < cols; i++) {
        max_val = fmaxf(max_val, row_input[i]);
    }

    // 计算 exp(x - max) 并求和
    float sum = 0.0f;
    for (int i = 0; i < cols; i++) {
        row_output[i] = expf(row_input[i] - max_val);
        sum += row_output[i];
    }

    // 归一化
    for (int i = 0; i < cols; i++) {
        row_output[i] /= sum;
    }
}

torch::Tensor fused_softmax_cuda(torch::Tensor input) {
    auto output = torch::zeros_like(input);
    int rows = input.size(0);
    int cols = input.size(1);

    fused_softmax_kernel<<<rows, 1>>>(
        input.data_ptr<float>(),
        output.data_ptr<float>(),
        rows,
        cols
    );

    return output;
}
"""

cpp_source_cuda = """
#include <torch/extension.h>
torch::Tensor fused_softmax_cuda(torch::Tensor input);

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("fused_softmax", &fused_softmax_cuda, "Fused Softmax CUDA");
}
"""

# 只在存在 CUDA 时才编译
if torch.cuda.is_available():
    from torch.utils.cpp_extension import load_inline
    fused_softmax_ext = load_inline(
        name="fused_softmax",
        cpp_sources=cpp_source_cuda,
        cuda_sources=cuda_source,
        functions=["fused_softmax"],
        verbose=True
    )

    # 测试
    x = torch.randn(4, 8, device="cuda")
    result = fused_softmax_ext.fused_softmax(x)
    expected = torch.softmax(x, dim=1)
    print(f"Max difference: {(result - expected).abs().max().item():.6f}")

使用 setup.py 构建包

# setup.py
from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension

setup(
    name="custom_ops",
    ext_modules=[
        CUDAExtension(
            name="custom_ops",
            sources=[
                "custom_ops/ops.cpp",
                "custom_ops/ops_cuda.cu",
            ],
            extra_compile_args={
                "cxx": ["-O3"],
                "nvcc": ["-O3", "--use_fast_math"],
            }
        )
    ],
    cmdclass={
        "build_ext": BuildExtension
    }
)
# 构建：python setup.py install

4. 内存优化技巧

GPU 内存性能分析

import torch

def print_gpu_memory_stats():
    """输出 GPU 内存统计信息的工具函数"""
    if torch.cuda.is_available():
        print(f"Allocated: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
        print(f"Reserved:  {torch.cuda.memory_reserved() / 1024**3:.2f} GB")
        print(f"Max Allocated: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")
        print(torch.cuda.memory_summary(abbreviated=True))

# 开始追踪内存
torch.cuda.reset_peak_memory_stats()
print_gpu_memory_stats()

# 加载模型
model = TransformerBlock().cuda()
print_gpu_memory_stats()

Gradient Checkpointing（Activation Recomputation）

Gradient Checkpointing 在前向传播时不保存中间激活值，而是在反向传播需要时再重新计算。它能大幅节省内存，但计算时间会增加约 30%。

from torch.utils.checkpoint import checkpoint, checkpoint_sequential
import torch.nn as nn

class DeepTransformer(nn.Module):
    def __init__(self, num_layers=12, d_model=512, nhead=8):
        super().__init__()
        self.layers = nn.ModuleList([
            TransformerBlock(d_model, nhead) for _ in range(num_layers)
        ])

    def forward_with_checkpointing(self, x):
        """为每一层应用 gradient checkpointing"""
        for layer in self.layers:
            # checkpoint 不保存中间激活值
            x = checkpoint(layer, x, use_reentrant=False)
        return x

    def forward_sequential_checkpointing(self, x):
        """为 sequential 模块应用 checkpointing"""
        # 按 4 层为一组分段
        x = checkpoint_sequential(self.layers, segments=3, input=x)
        return x

# 内存比较
model = DeepTransformer(num_layers=24).cuda()
x = torch.randn(16, 512, 512, device="cuda")

# 常规 forward
torch.cuda.reset_peak_memory_stats()
out = model(x) if hasattr(model, 'forward') else model.forward_with_checkpointing(x)
normal_mem = torch.cuda.max_memory_allocated()

# checkpointing forward
torch.cuda.reset_peak_memory_stats()
out = model.forward_with_checkpointing(x)
checkpoint_mem = torch.cuda.max_memory_allocated()

print(f"Normal memory: {normal_mem / 1024**3:.2f} GB")
print(f"Checkpoint memory: {checkpoint_mem / 1024**3:.2f} GB")
print(f"Memory saved: {(normal_mem - checkpoint_mem) / 1024**3:.2f} GB")

Gradient Accumulation

def train_with_gradient_accumulation(
    model, optimizer, dataloader, criterion,
    accumulation_steps=4
):
    """
    用较小的 GPU 内存实现较大的有效批大小（effective batch size）
    """
    model.train()
    optimizer.zero_grad()
    total_loss = 0

    for batch_idx, (data, target) in enumerate(dataloader):
        data, target = data.cuda(), target.cuda()

        # Forward pass
        output = model(data)
        # 按 accumulation_steps 缩放损失
        loss = criterion(output, target) / accumulation_steps
        loss.backward()

        total_loss += loss.item() * accumulation_steps

        # 每 accumulation_steps 步更新一次
        if (batch_idx + 1) % accumulation_steps == 0:
            # 梯度裁剪（可选）
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()
            optimizer.zero_grad()

    return total_loss / len(dataloader)

Mixed Precision Training（AMP）

from torch.cuda.amp import autocast, GradScaler

def train_with_amp(model, optimizer, dataloader, criterion):
    """用 Automatic Mixed Precision 节省内存 + 提升速度"""
    scaler = GradScaler()
    model.train()

    for data, target in dataloader:
        data, target = data.cuda(), target.cuda()
        optimizer.zero_grad()

        # 在 autocast 上下文中执行 float16 运算
        with autocast():
            output = model(data)
            loss = criterion(output, target)

        # 用缩放后的梯度做 backward
        scaler.scale(loss).backward()

        # 反缩放后再做梯度裁剪
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)

        # Optimizer 步进（含 NaN/Inf 检查）
        scaler.step(optimizer)
        scaler.update()

8-bit Optimizer

# 需要 bitsandbytes 库：pip install bitsandbytes
try:
    import bitsandbytes as bnb

    # 用 8-bit Adam 代替常规 Adam
    optimizer_8bit = bnb.optim.Adam8bit(
        model.parameters(),
        lr=1e-4,
        betas=(0.9, 0.999)
    )

    # PagedAdam（可卸载到 CPU）
    optimizer_paged = bnb.optim.PagedAdam(
        model.parameters(),
        lr=1e-4
    )

    print("8-bit optimizer loaded successfully")
except ImportError:
    print("bitsandbytes not installed, using regular Adam")
    optimizer_8bit = torch.optim.Adam(model.parameters(), lr=1e-4)

5. functorch 与 vmap

vmap - 批量化运算

vmap 能把作用于单个样本的函数，高效地应用到整个批次上。

import torch
from torch import vmap

# 作用于单个样本的函数
def single_linear(weight, bias, x):
    return weight @ x + bias

# 用 vmap 做批量处理
batched_linear = vmap(single_linear)

# 批量数据
batch_size = 32
weight = torch.randn(batch_size, 10, 5)
bias = torch.randn(batch_size, 10)
x = torch.randn(batch_size, 5)

# 自动完成批量处理
result = batched_linear(weight, bias, x)
print(f"Result shape: {result.shape}")  # (32, 10)

grad - 函数式梯度

from torch.func import grad, vmap, functional_call

# 函数式梯度
def scalar_loss(params, x, y):
    pred = functional_call(model, params, (x,))
    return ((pred - y) ** 2).mean()

# 对参数求梯度
params = dict(model.named_parameters())
grad_fn = grad(scalar_loss)

x = torch.randn(1, 10)
y = torch.randn(1, 5)
grads = grad_fn(params, x, y)
print({k: v.shape for k, v in grads.items()})

集成模型（Ensemble）- 使用 vmap

from torch.func import stack_module_state, functional_call, vmap

def create_ensemble(model_class, num_models, *args, **kwargs):
    """借助 vmap 实现的高效集成"""
    models = [model_class(*args, **kwargs) for _ in range(num_models)]

    # 把所有模型的参数堆叠起来
    params, buffers = stack_module_state(models)

    # 单个模型的 forward
    base_model = model_class(*args, **kwargs)

    def single_forward(params, buffers, x):
        return functional_call(base_model, (params, buffers), (x,))

    # 用 vmap 并行执行所有模型
    ensemble_forward = vmap(single_forward, in_dims=(0, 0, None))

    return ensemble_forward, params, buffers

# 使用示例
ensemble_fn, params, buffers = create_ensemble(
    nn.Linear, num_models=5, in_features=10, out_features=5
)

x = torch.randn(32, 10)
ensemble_out = ensemble_fn(params, buffers, x)
print(f"Ensemble output shape: {ensemble_out.shape}")  # (5, 32, 5)

元学习（MAML）- 结合 grad 与 vmap

from torch.func import grad, vmap, functional_call

def inner_loop(params, support_x, support_y, base_model, lr=0.01, steps=5):
    """MAML 内循环"""
    adapted_params = {k: v.clone() for k, v in params.items()}

    for _ in range(steps):
        def loss_fn(params):
            pred = functional_call(base_model, params, (support_x,))
            return ((pred - support_y) ** 2).mean()

        grads = grad(loss_fn)(adapted_params)
        adapted_params = {
            k: p - lr * grads[k]
            for k, p in adapted_params.items()
        }

    return adapted_params

6. PyTorch Profiler

基础性能分析

import torch
from torch.profiler import profile, record_function, ProfilerActivity

model = TransformerBlock().cuda()
x = torch.randn(32, 128, 512, device="cuda")

# 运行 profiler
with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    record_shapes=True,
    profile_memory=True,
    with_stack=True
) as prof:
    with record_function("model_inference"):
        for _ in range(10):
            output = model(x)

# 输出结果
print(prof.key_averages().table(
    sort_by="cuda_time_total",
    row_limit=20
))

# 导出 Chrome Trace
prof.export_chrome_trace("trace.json")
# 在 chrome://tracing 中打开

与 TensorBoard 集成

from torch.profiler import profile, ProfilerActivity, tensorboard_trace_handler

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    schedule=torch.profiler.schedule(
        wait=1,    # 前 1 步等待
        warmup=1,  # 1 步预热
        active=3,  # 3 步进行分析
        repeat=2   # 重复 2 次
    ),
    on_trace_ready=tensorboard_trace_handler("./log/profiler"),
    record_shapes=True,
    profile_memory=True
) as prof:
    for step, (data, target) in enumerate(dataloader):
        train_step(model, optimizer, data, target)
        prof.step()  # 按调度进行分析

# tensorboard --logdir=./log/profiler

详细内存分析

# 内存快照（PyTorch 2.1+）
torch.cuda.memory._record_memory_history(max_entries=100000)

# 执行代码
x = torch.randn(100, 100, device="cuda")
y = x @ x.T
z = y.sum()

# 保存快照
snapshot = torch.cuda.memory._snapshot()
torch.cuda.memory._dump_snapshot("memory_snapshot.pickle")

# 分析
print(f"Active allocations: {len(snapshot['segments'])}")

7. TorchScript

torch.jit.script 与 torch.jit.trace

import torch
import torch.nn as nn

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

    def forward(self, x, flag: bool = True):
        if flag:  # 有条件分支时无法用 trace
            return torch.relu(self.linear(x))
        else:
            return torch.sigmoid(self.linear(x))

model = ConditionalModel()

# script：包含控制流（推荐）
scripted_model = torch.jit.script(model)
print(scripted_model.code)

# trace：只捕获单一执行路径
x = torch.randn(1, 10)
traced_model = torch.jit.trace(model, x)
# 注意：flag=False 的分支不会被捕获

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

TorchScript 优化

# 应用优化 pass
scripted = torch.jit.script(model)
optimized = torch.jit.optimize_for_inference(scripted)

# 导出到 C++ 环境
# 在 C++ 中：torch::jit::script::Module m = torch::jit::load("model.pt");

8. 动态形状与 torch.export

使用 torch.export

import torch
from torch.export import export

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

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

model = MyModel()
example_inputs = (torch.randn(2, 10),)

# 指定动态形状
from torch.export import Dim
batch = Dim("batch", min=1, max=100)

# 导出模型
exported = export(
    model,
    example_inputs,
    dynamic_shapes={"x": {0: batch}}
)

print(exported)
print(exported.graph_module.code)

# 运行 ExportedProgram
result = exported.module()(torch.randn(5, 10))
print(f"Result shape: {result.shape}")

9. 自定义 Dataset 与 Sampler

IterableDataset

from torch.utils.data import IterableDataset, DataLoader
import torch

class StreamingDataset(IterableDataset):
    """以流式方式处理大规模数据"""

    def __init__(self, data_paths, transform=None):
        self.data_paths = data_paths
        self.transform = transform

    def __iter__(self):
        worker_info = torch.utils.data.get_worker_info()

        if worker_info is None:
            # 单进程
            paths = self.data_paths
        else:
            # 多进程：切分数据
            per_worker = len(self.data_paths) // worker_info.num_workers
            start = worker_info.id * per_worker
            end = start + per_worker
            paths = self.data_paths[start:end]

        for path in paths:
            # 从文件流式读取数据
            data = self._load_file(path)
            for sample in data:
                if self.transform:
                    sample = self.transform(sample)
                yield sample

    def _load_file(self, path):
        # 实际实现中应从文件加载
        return [torch.randn(10) for _ in range(100)]

# 在 DataLoader 中使用 persistent_workers
dataset = StreamingDataset(data_paths=["file1.pt", "file2.pt"])
dataloader = DataLoader(
    dataset,
    batch_size=32,
    num_workers=4,
    persistent_workers=True,  # 复用工作进程
    pin_memory=True,           # 优化 GPU 传输
    prefetch_factor=2          # 预取的批次数
)

自定义 Sampler

from torch.utils.data import Sampler
import numpy as np

class BalancedClassSampler(Sampler):
    """用于解决类别不平衡问题的加权采样"""

    def __init__(self, dataset, num_samples_per_class=None):
        self.dataset = dataset
        labels = [dataset[i][1] for i in range(len(dataset))]
        self.labels = torch.tensor(labels)

        # 按类别记录索引
        self.class_indices = {}
        for cls in torch.unique(self.labels):
            self.class_indices[cls.item()] = (
                self.labels == cls
            ).nonzero(as_tuple=True)[0].tolist()

        self.num_classes = len(self.class_indices)
        self.num_samples_per_class = (
            num_samples_per_class or
            max(len(v) for v in self.class_indices.values())
        )

    def __iter__(self):
        indices = []
        for cls_idx in self.class_indices.values():
            # 从每个类别中抽取相同数量的样本（有放回）
            sampled = np.random.choice(
                cls_idx,
                self.num_samples_per_class,
                replace=True
            ).tolist()
            indices.extend(sampled)

        # 打乱顺序
        np.random.shuffle(indices)
        return iter(indices)

    def __len__(self):
        return self.num_classes * self.num_samples_per_class

10. PyTorch Lightning

LightningModule 完整示例

import pytorch_lightning as pl
from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping
from pytorch_lightning.loggers import TensorBoardLogger
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset

class LightningTransformer(pl.LightningModule):
    def __init__(
        self,
        d_model=512,
        nhead=8,
        num_layers=6,
        num_classes=10,
        learning_rate=1e-4
    ):
        super().__init__()
        self.save_hyperparameters()  # 自动保存 HParams

        self.encoder = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(
                d_model=d_model,
                nhead=nhead,
                batch_first=True
            ),
            num_layers=num_layers
        )
        self.classifier = nn.Linear(d_model, num_classes)
        self.criterion = nn.CrossEntropyLoss()

    def forward(self, x):
        encoded = self.encoder(x)
        # 对序列做平均池化
        pooled = encoded.mean(dim=1)
        return self.classifier(pooled)

    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = self.criterion(logits, y)
        acc = (logits.argmax(dim=1) == y).float().mean()

        # 日志记录
        self.log("train_loss", loss, prog_bar=True)
        self.log("train_acc", acc, prog_bar=True)
        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = self.criterion(logits, y)
        acc = (logits.argmax(dim=1) == y).float().mean()

        self.log("val_loss", loss, prog_bar=True)
        self.log("val_acc", acc, prog_bar=True)

    def configure_optimizers(self):
        optimizer = torch.optim.AdamW(
            self.parameters(),
            lr=self.hparams.learning_rate,
            weight_decay=0.01
        )
        # 余弦退火调度器
        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
            optimizer, T_max=100
        )
        return {
            "optimizer": optimizer,
            "lr_scheduler": {
                "scheduler": scheduler,
                "monitor": "val_loss"
            }
        }


class LightningDataModule(pl.LightningDataModule):
    def __init__(self, batch_size=32):
        super().__init__()
        self.batch_size = batch_size

    def setup(self, stage=None):
        # 虚拟数据
        x = torch.randn(1000, 32, 512)
        y = torch.randint(0, 10, (1000,))
        dataset = TensorDataset(x, y)
        n = len(dataset)
        self.train_ds = torch.utils.data.Subset(dataset, range(int(0.8*n)))
        self.val_ds = torch.utils.data.Subset(dataset, range(int(0.8*n), n))

    def train_dataloader(self):
        return DataLoader(self.train_ds, batch_size=self.batch_size, shuffle=True)

    def val_dataloader(self):
        return DataLoader(self.val_ds, batch_size=self.batch_size)


# 执行训练
model = LightningTransformer()
dm = LightningDataModule()

# 配置回调
callbacks = [
    ModelCheckpoint(
        monitor="val_loss",
        save_top_k=3,
        mode="min",
        filename="transformer-{epoch:02d}-{val_loss:.2f}"
    ),
    EarlyStopping(monitor="val_loss", patience=10, mode="min")
]

# Trainer
trainer = pl.Trainer(
    max_epochs=100,
    accelerator="auto",    # 自动检测 GPU
    devices="auto",
    callbacks=callbacks,
    logger=TensorBoardLogger("tb_logs", name="transformer"),
    gradient_clip_val=1.0,  # 梯度裁剪
    accumulate_grad_batches=4,  # Gradient accumulation
    precision="16-mixed",   # AMP
)

trainer.fit(model, dm)

11. 模型量化（Quantization）

动态量化（Dynamic Quantization）

import torch
import torch.nn as nn
from torch.ao.quantization import quantize_dynamic

# 仅用于推理的量化 - 最简单的方法
model = nn.Sequential(
    nn.Linear(512, 256),
    nn.ReLU(),
    nn.Linear(256, 128)
)
model.eval()

# 把 Linear 和 LSTM 层量化为 int8
quantized_model = quantize_dynamic(
    model,
    {nn.Linear, nn.LSTM},  # 要量化的层类型
    dtype=torch.qint8
)

# 体积比较
import os
torch.save(model.state_dict(), "model_fp32.pt")
torch.save(quantized_model.state_dict(), "model_int8.pt")
fp32_size = os.path.getsize("model_fp32.pt")
int8_size = os.path.getsize("model_int8.pt")
print(f"FP32 size: {fp32_size / 1024:.1f} KB")
print(f"INT8 size: {int8_size / 1024:.1f} KB")
print(f"Compression: {fp32_size / int8_size:.1f}x")

静态量化（Static Quantization）

import torch
from torch.ao.quantization import (
    get_default_qconfig,
    prepare,
    convert,
    QConfig
)

class QuantizableModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.quant = torch.ao.quantization.QuantStub()
        self.linear1 = nn.Linear(512, 256)
        self.relu = nn.ReLU()
        self.linear2 = nn.Linear(256, 10)
        self.dequant = torch.ao.quantization.DeQuantStub()

    def forward(self, x):
        x = self.quant(x)
        x = self.linear1(x)
        x = self.relu(x)
        x = self.linear2(x)
        x = self.dequant(x)
        return x

model = QuantizableModel()
model.eval()

# 量化配置
model.qconfig = get_default_qconfig("fbgemm")  # 用于 x86 CPU

# 准备（插入 observer）
prepared_model = prepare(model)

# 用校准数据收集统计信息
with torch.no_grad():
    for _ in range(100):
        calibration_data = torch.randn(32, 512)
        prepared_model(calibration_data)

# 转换为量化模型
quantized_static = convert(prepared_model)

# 推理
x = torch.randn(1, 512)
with torch.no_grad():
    output = quantized_static(x)
print(f"Output shape: {output.shape}")

QAT（Quantization-Aware Training）

from torch.ao.quantization import (
    get_default_qat_qconfig,
    prepare_qat,
    convert
)

model = QuantizableModel()
model.train()

# QAT 配置
model.qconfig = get_default_qat_qconfig("fbgemm")

# 准备 QAT（插入伪量化节点）
prepared_qat = prepare_qat(model)

# 训练（把量化误差纳入训练过程）
optimizer = torch.optim.SGD(prepared_qat.parameters(), lr=0.0001)
for epoch in range(10):
    for x, y in dummy_dataloader():
        output = prepared_qat(x)
        loss = nn.functional.cross_entropy(output, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

# 训练完成后转换
prepared_qat.eval()
quantized_qat = convert(prepared_qat)


def dummy_dataloader():
    for _ in range(10):
        yield torch.randn(32, 512), torch.randint(0, 10, (32,))

12. Tensor Parallelism 与 Pipeline Parallelism

DeviceMesh 与 DTensor API

import torch
import torch.distributed as dist
from torch.distributed.tensor.parallel import (
    parallelize_module,
    ColwiseParallel,
    RowwiseParallel
)
from torch.distributed._tensor import DeviceMesh

# 初始化分布式训练
def setup_distributed():
    dist.init_process_group(backend="nccl")
    torch.cuda.set_device(dist.get_rank())

class TransformerMLP(nn.Module):
    def __init__(self, d_model=1024, dim_feedforward=4096):
        super().__init__()
        self.fc1 = nn.Linear(d_model, dim_feedforward)
        self.act = nn.GELU()
        self.fc2 = nn.Linear(dim_feedforward, d_model)

    def forward(self, x):
        return self.fc2(self.act(self.fc1(x)))

# 应用 Tensor Parallelism
def apply_tensor_parallel(model, mesh):
    """
    fc1 按列（Column-wise）切分（切分输出维度）
    fc2 按行（Row-wise）切分（切分输入维度）
    """
    parallelize_module(
        model,
        mesh,
        {
            "fc1": ColwiseParallel(),
            "fc2": RowwiseParallel(),
        }
    )
    return model

# 运行示例（2 GPU 环境）
# device_mesh = DeviceMesh("cuda", [0, 1])
# model = TransformerMLP()
# model = apply_tensor_parallel(model, device_mesh)

FSDP（Fully Sharded Data Parallel）

from torch.distributed.fsdp import (
    FullyShardedDataParallel as FSDP,
    MixedPrecision,
    BackwardPrefetch,
    ShardingStrategy
)
from torch.distributed.fsdp.wrap import (
    size_based_auto_wrap_policy,
    transformer_auto_wrap_policy
)
import functools

def setup_fsdp_model(model):
    """FSDP 配置 - 把大规模模型分散到多张 GPU 上"""

    # 混合精度配置
    mp_policy = MixedPrecision(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.float32,
        buffer_dtype=torch.bfloat16,
    )

    # 以 TransformerBlock 为单位进行包装
    auto_wrap_policy = functools.partial(
        transformer_auto_wrap_policy,
        transformer_layer_cls={TransformerBlock}
    )

    model = FSDP(
        model,
        auto_wrap_policy=auto_wrap_policy,
        mixed_precision=mp_policy,
        sharding_strategy=ShardingStrategy.FULL_SHARD,
        backward_prefetch=BackwardPrefetch.BACKWARD_PRE,
        device_id=torch.cuda.current_device()
    )

    return model

结语

本指南梳理了 PyTorch 的高级技巧：

torch.compile：无需修改代码，即可实现 2 倍以上的性能提升
Custom Autograd：实现特殊的梯度计算
CUDA Extensions：把 GPU 内核整合进 PyTorch
内存优化：Gradient Checkpointing、AMP、8-bit Optimizer
functorch/vmap：用函数式 API 实现批量处理和元学习
PyTorch Profiler：分析性能瓶颈
TorchScript/export：部署优化
PyTorch Lightning：代码结构化与训练自动化
量化：用 INT8 压缩模型体积
分布式训练：Tensor Parallel、FSDP

各项技巧相辅相成，实际项目中通常会组合使用多种技巧。尤其在大规模模型训练时，AMP + Gradient Checkpointing + FSDP + torch.compile 的组合威力强大。