PyTorch 高级技巧
🪝 钩子函数（Hooks）
Hooks 让你在不修改 forward() 代码的情况下，在模型的前向/反向传播中插入自定义逻辑。非常适合调试、特征提取、梯度分析和可视化。
前向钩子（Forward Hook）
python
# 注册前向钩子
def forward_hook(module, input, output):
"""module: 被 hook 的层
input: 输入元组
output: 输出张量"""
print(f'{module.__class__.__name__}: input shape {input[0].shape}')
print(f' → output shape {output.shape}')
# 可以保存输出用于特征可视化
# 注册到某一层
handle = model.layer3.register_forward_hook(forward_hook)
# 用完记得移除
handle.remove()反向钩子（Backward Hook）
python
# 查看梯度统计
def backward_hook(module, grad_input, grad_output):
"""检查梯度是否正常"""
for i, g in enumerate(grad_output):
if g is not None:
print(f' grad_output[{i}]: mean={g.mean():.6f}, std={g.std():.6f}, '
f'nan={torch.isnan(g).any()}, inf={torch.isinf(g).any()}')
handle = model.fc2.register_full_backward_hook(backward_hook)参数钩子
python
# 监控参数更新
def param_hook(grad):
return grad * 0.5 # 缩放梯度
model.fc1.weight.register_hook(param_hook)实战：特征提取
python
# 提取中间层特征（不改模型代码）
features = {}
def get_features(name):
def hook(model, input, output):
features[name] = output.detach()
return hook
model.layer1.register_forward_hook(get_features('layer1'))
model.layer4.register_forward_hook(get_features('layer4'))
# 前向传播后
output = model(input_image)
print(features['layer1'].shape) # 浅层特征
print(features['layer4'].shape) # 深层特征
# 可视化特征图
import matplotlib.pyplot as plt
feat = features['layer1'][0] # batch 第一个样本
fig, axes = plt.subplots(4, 8, figsize=(12, 6))
for i, ax in enumerate(axes.flat):
if i < feat.shape[0]:
ax.imshow(feat[i].cpu(), cmap='viridis')
ax.axis('off')实战：梯度裁剪 + 监控
python
def grad_clip_and_log(module, grad_input, grad_output):
"""在反向传播时自动裁剪并记录梯度"""
for g in grad_output:
if g is not None:
torch.nn.utils.clip_grad_norm_(g, max_norm=1.0)
print(f'Grad norm: {g.norm():.4f}')
handle = model.register_full_backward_hook(grad_clip_and_log)🔍 性能分析（Profiling）
使用 PyTorch Profiler
python
from torch.profiler import profile, record_function, ProfilerActivity
with profile(
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1),
on_trace_ready=torch.profiler.tensorboard_trace_handler('./log/profiler'),
record_shapes=True,
profile_memory=True,
with_stack=True
) as prof:
for step, (X, y) in enumerate(train_loader):
with record_function("forward"):
output = model(X)
loss = loss_fn(output, y)
with record_function("backward"):
loss.backward()
with record_function("update"):
optimizer.step()
prof.step()
if step >= 5:
break
# 查看结果
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
# 在 TensorBoard 中查看
# tensorboard --logdir=./log/profiler简单计时
python
import time
# CPU 计时
start = time.time()
output = model(input)
print(f'Inference: {time.time() - start:.4f}s')
# GPU 同步计时（更准确）
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record()
output = model(input)
end.record()
torch.cuda.synchronize()
print(f'GPU Inference: {start.elapsed_time(end):.2f}ms')单层耗时分析
python
# 统计每层的 FLOPs 和参数量
# pip install fvcore
from fvcore.nn import FlopCountAnalysis, parameter_count
input_tensor = torch.randn(1, 3, 224, 224)
flops = FlopCountAnalysis(model, input_tensor)
params = parameter_count(model)
print(f'Total FLOPs: {flops.total() / 1e9:.2f} G')
print(f'Total Params: {params[""] / 1e6:.2f} M')🔢 模型量化（Quantization）
量化将模型从 FP32 降到 INT8，大幅减少模型体积和推理延迟。
动态量化（最简单）
python
# 只量化 Linear 和 LSTM 层的权重（激活保持 FP32）
quantized_model = torch.quantization.quantize_dynamic(
model, # 原始模型
{nn.Linear, nn.LSTM}, # 要量化的层类型
dtype=torch.qint8 # 量化目标类型
)
# 适用于 LSTM、Transformer（对 CNN 效果一般）静态量化（更高效）
python
# 1. 设置量化配置
model.qconfig = torch.ao.quantization.get_default_qconfig('x86')
# 或 'fbgemm'（x86 CPU），'qnnpack'（ARM CPU）
# 2. 插入观察器
model_prepared = torch.ao.quantization.prepare(model)
# 3. 用校准数据运行（收集激活分布）
with torch.no_grad():
for X, _ in calibration_loader:
model_prepared(X)
# 4. 转换为量化模型
model_quantized = torch.ao.quantization.convert(model_prepared)PyTorch 2.0 torch.compile 量化（推荐）
python
# 更简单：torch.compile 内置量化优化
model = torch.compile(model, mode='reduce-overhead')
# 或
from torch.ao.quantization import quantize_pt2e
model = quantize_pt2e(model, ...) # 当前实验性🔌 自定义算子（Custom C++/CUDA Extension）
Python 层面（torch.autograd.Function）
python
class MyReLU(torch.autograd.Function):
@staticmethod
def forward(ctx, input):
"""前向传播"""
ctx.save_for_backward(input) # 保存用于反向传播
return input.clamp(min=0)
@staticmethod
def backward(ctx, grad_output):
"""反向传播"""
input, = ctx.saved_tensors
grad_input = grad_output.clone()
grad_input[input < 0] = 0
return grad_input
# 使用
relu_custom = MyReLU.apply
x = torch.randn(5, requires_grad=True)
y = relu_custom(x)示例: GELU 自定义实现
python
class GELU(torch.autograd.Function):
@staticmethod
def forward(ctx, x):
ctx.save_for_backward(x)
return 0.5 * x * (1 + torch.tanh(
(2 / torch.pi) ** 0.5 * (x + 0.044715 * x ** 3)))
@staticmethod
def backward(ctx, grad_output):
x, = ctx.saved_tensors
# GELU 的导数
cdf = 0.5 * (1 + torch.erf(x / (2 ** 0.5)))
pdf = torch.exp(-0.5 * x ** 2) / ((2 * torch.pi) ** 0.5)
return grad_output * (cdf + x * pdf)
gelu = GELU.apply📜 TorchScript --- 序列化与跨语言部署
Tracing（追踪）vs Scripting（脚本）
python
# Tracing: 喂入样例输入，记录执行路径（不支持控制流）
class SimpleModel(nn.Module):
def forward(self, x):
return x * 2 + 1
model = SimpleModel()
example = torch.randn(1, 3)
traced = torch.jit.trace(model, example)
# Scripting: 分析 Python 代码（支持控制流）
@torch.jit.script
def my_loop(x):
for i in range(10):
x = x + i
return x
class ModelWithControl(nn.Module):
def forward(self, x, n):
if n > 0: # ← 控制流
x = x * 2
return x
scripted = torch.jit.script(ModelWithControl())保存与加载
python
# 保存
traced.save('model_traced.pt')
# 加载（无需 Python 代码！）
model = torch.jit.load('model_traced.pt')
# 在 C++ 中使用
# module = torch::jit::load("model_traced.pt");🏗️ 参数共享与绑定
python
# 让两个层共享参数（权重绑定）
class SharedModel(nn.Module):
def __init__(self):
super().__init__()
self.encoder = nn.Linear(100, 50)
self.decoder = nn.Linear(50, 100)
# 绑定 decoder 权重 = encoder 权重的转置
self.decoder.weight = self.encoder.weight # 共享参数！
model = SharedModel()
# encoder.weight 和 decoder.weight 是同一个张量🧪 torch.inference_mode() vs torch.no_grad()
python
# torch.no_grad(): 禁用梯度计算，但仍追踪 view/reshape
with torch.no_grad():
y = model(x)
# torch.inference_mode(): 更激进，完全禁用 autograd（推荐推理用）
with torch.inference_mode():
y = model(x)
# 区别: inference_mode 更快、更省内存，但不能同时和 autograd 混用📐 动态 Padding 批处理
python
# 处理变长输入的高效方案
from torch.nn.utils.rnn import pad_sequence, pack_padded_sequence, pad_packed_sequence
# 变长序列
seqs = [torch.randn(3, 128), torch.randn(5, 128), torch.randn(2, 128)]
lengths = torch.tensor([3, 5, 2])
# Padding
padded = pad_sequence(seqs, batch_first=True) # (3, 5, 128)
# 打包（跳过 padding 的计算）
packed = pack_padded_sequence(padded, lengths, batch_first=True, enforce_sorted=False)
# RNN 前向传播
output_packed, (h, c) = lstm(packed)
# 解包
output, _ = pad_packed_sequence(output_packed, batch_first=True)🗑️ 自定义 nn.Module 的高级特性
python
class AdvancedModule(nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(10, 10)
# 非参数缓冲区（随模型移动，但不参与训练）
self.register_buffer('running_mean', torch.zeros(10))
# 常量（不会被 state_dict 保存）
self.register_buffer('const', torch.tensor([1.0]), persistent=False)
def extra_repr(self):
"""自定义 print(model) 的输出"""
return 'special_config=42'
def train(self, mode=True):
"""自定义训练/评估模式切换"""
super().train(mode)
print(f'Switched to {"train" if mode else "eval"} mode')🧩 nn.ModuleList 和 nn.ModuleDict
python
# ModuleList: 按索引访问（类似 Python list）
layers = nn.ModuleList([
nn.Linear(10, 20) for _ in range(5)
])
x = layers[0](x) # 可以索引访问
x = layers[1](x)
# ModuleDict: 按键名访问（类似 Python dict）
blocks = nn.ModuleDict({
'encoder': nn.Linear(10, 5),
'decoder': nn.Linear(5, 10),
})
x = blocks['encoder'](x)
# 关键: 它们的参数会被正确注册，optimizer 和 .to() 都会生效
# 而普通的 Python list/dict 里的 Module 不会自动注册！🎛️ 可微分数据增强
python
# kornia: 在 GPU 上可微分的数据增强
# pip install kornia
import kornia.augmentation as K
augs = nn.Sequential(
K.RandomHorizontalFlip(p=0.5),
K.RandomRotation(degrees=15.0),
K.ColorJitter(0.2, 0.2, 0.2, 0.1),
)
# 直接在训练循环中 GPU 上执行（比 torchvision transforms 快得多）
x_aug = augs(x) # x 在 GPU 上，结果也在 GPU 上🔧 调试技巧
python
# 1. 检测 NaN/Inf
if torch.isnan(loss):
print('NaN detected!')
break
if torch.isinf(loss):
print('Inf detected!')
break
# 2. 异常检测（开发时开启，会明显变慢）
torch.autograd.set_detect_anomaly(True)
# 3. 打印所有参数的梯度范数
total_norm = 0.0
for p in model.parameters():
if p.grad is not None:
total_norm += p.grad.data.norm(2).item() ** 2
print(f'Gradient norm: {total_norm ** 0.5:.4f}')
# 4. 设置断点调试反向传播
with torch.autograd.profiler.emit_nvtx():
loss.backward() # 可在 Nsight Systems 中查看
# 5. 查看每层输出的统计
for name, module in model.named_modules():
def hook(name):
def fn(module, input, output):
if isinstance(output, torch.Tensor):
print(f'{name}: {output.mean():.4f} ± {output.std():.4f}')
return fn
module.register_forward_hook(hook(name))📝 速查表
| 需求 | 代码 |
|---|---|
| 前向钩子 | module.register_forward_hook(fn) |
| 反向钩子 | module.register_full_backward_hook(fn) |
| 推理模式 | with torch.inference_mode(): |
| Profiler | with profile(...) as prof: |
| 动态量化 | torch.quantization.quantize_dynamic(model, ...) |
| 自定义算子 | class MyFunc(torch.autograd.Function): |
| Tracing | torch.jit.trace(model, example) |
| Scripting | torch.jit.script(model) |
| 注册 buffer | self.register_buffer('name', tensor) |
| ModuleList | nn.ModuleList([...]) |
| NaN 检测 | torch.isnan(tensor) |
| 异常检测 | torch.autograd.set_detect_anomaly(True) |
| 参数共享 | 直接赋值 .weight = other.weight |
\[pytorch-总览\|← 返回总览\]