torch

PyTorch 常用速查

安装

pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cu128

基础配置

环境初始化 (随机种子)

固定随机种子以保证实验可复现性。

import random
import os
import numpy as np
import torch

def seed_everything(seed):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True

设备选择

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

张量操作

常用创建

x = torch.zeros(2, 3)       # 全0
x = torch.ones(2, 3)        # 全1
x = torch.randn(2, 3)       # 标准正态分布
x = torch.tensor([1, 2, 3]) # 从list创建
x = torch.from_numpy(np_arr)# 从numpy创建

维度变换

x.view(3, 2)        # 改变形状 (需内存连续)
x.reshape(3, 2)     # 改变形状 (通用)
x.permute(1, 0)     # 交换维度
x.unsqueeze(0)      # 增加维度
x.squeeze(0)        # 减少维度

梯度控制

# 开启梯度，带下划线'_'，表示原地修改。
x.requires_grad_(True)

# 分离梯度 (常用于生成伪标签或截断反向传播)
with torch.no_grad():
    y = model(x)
    y = y.detach()

模型与训练

训练循环模板

model.train()
optimizer.zero_grad()       # 1. 梯度清零
output = model(data)        # 2. 前向传播
loss = criterion(output, target)
loss.backward()             # 3. 反向传播
optimizer.step()            # 4. 参数更新

Note: optimizer.zero_grad() 必不可少，否则梯度会累加（相当于变相增加 batch size），可能导致显存溢出或训练错误。

推理模式

model.eval()
with torch.no_grad():
    # 不计算梯度，节省显存和计算资源
    output = model(data)

获取中间层输出 (Hook)

通过 Hook 机制提取模型中间层的特征。

def get_activation(mem, name):
    def get_output_hook(module, input, output):
        mem[name] = output
    return get_output_hook

def add_hook(net, mem, mapping_layers):
    for n, m in net.named_modules():
        if n in mapping_layers:
            m.register_forward_hook(get_activation(mem, n))

# 使用示例
acts = {}
target_layers = ['layer1', 'layer2.conv1']
add_hook(model, acts, target_layers)
# inference后，acts中将包含对应层的输出

显存管理

手动清理

当遇到显存不足或需要强制释放资源时使用。

import gc
def cleanup():
    gc.collect()
    torch.cuda.empty_cache()
    # tcnn.free_temporary_memory() # 如果使用了tiny-cuda-nn

模型保存与加载

# 保存
torch.save(model.state_dict(), 'model.pth')

# 加载
model.load_state_dict(torch.load('model.pth'))