神经网络性能优化：识别和解决bottleneck的5种实战策略

1次阅读

共计 2796 个字符，预计需要花费 7 分钟才能阅读完成。

在实际项目中发现，当 GPU 利用率长期低于 50% 时，通常意味着存在严重的计算资源浪费。例如在 ResNet50 训练任务中，仅因数据加载和预处理阶段的 bottleneck，就可能导致 GPU 实际计算时间占比不到总训练时长的 40%。更糟糕的是，显存分配不当可能让 16GB 显存的显卡只能承载不到 8GB 的模型参数。

TensorRT 通过 kernel fusion（将多个操作合并为单个 CUDA 核函数）可以显著减少内存访问开销。典型场景如：将 Conv+BN+ReLU 三个独立操作融合为单个核函数。

import tensorrt as trt

# 创建 builder 和 network
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))

# 转换 PyTorch 模型到 ONNX
# 注意添加动态 shape 支持
dynamic_axes = {'input': {0: 'batch'}, 'output': {0: 'batch'}}
torch.onnx.export(model, dummy_input, "model.onnx", 
                  dynamic_axes=dynamic_axes)

# 使用 TensorRT 解析 ONNX
parser = trt.OnnxParser(network, logger)
with open("model.onnx", 'rb') as model:
    if not parser.parse(model.read()):
        for error in range(parser.num_errors):
            print(parser.get_error(error))
        raise ValueError("ONNX 解析失败")

# 配置优化参数
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB
profile = builder.create_optimization_profile()
profile.set_shape("input", (1,3,224,224), (8,3,224,224), (16,3,224,224))
config.add_optimization_profile(profile)

# 构建引擎
engine = builder.build_engine(network, config)

混合精度训练需要同时处理三件事：
– 使用 FP16 存储参数和计算梯度
– 维护 FP32 的主权重副本
– 动态损失缩放防止梯度下溢

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()  # 自动处理损失缩放

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

    with autocast():  # 自动选择 FP16/FP32
        output = model(data)
        loss = criterion(output, target)

    # 反向传播时自动缩放梯度
    scaler.scale(loss).backward()
    # 梯度更新时自动 unscale 并跳过溢出步骤
    scaler.step(optimizer)
    scaler.update()

通过 nsight systems 工具分析发现，当 CUDA 核心利用率超过 90% 而 Tensor Core 利用率低于 40% 时，应该：
– 将矩阵乘法的维度调整为 8 的倍数（Tensor Core 要求）
– 使用 torch.backends.cudnn.benchmark=True 自动寻找最优卷积算法

PyTorch profiler 是发现性能热点的利器：

with torch.profiler.profile(
    activities=[torch.profiler.ProfilerActivity.CPU,
                torch.profiler.ProfilerActivity.CUDA],
    schedule=torch.profiler.schedule(wait=1, warmup=1, active=3),
    on_trace_ready=torch.profiler.tensorboard_trace_handler('./log'),
    record_shapes=True
) as prof:
    for step, data in enumerate(dataloader):
        if step >= 5: break
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()
        prof.step()

关键指标解读：
– GPU Idle：GPU 空闲时间占比
– Kernel Time：各 CUDA 核函数执行时间
– Memcpy：内存拷贝耗时

不同代际 GPU 的优化策略差异：
– Volta 架构：需要显式启用 TF32
– Ampere 架构：支持稀疏张量运算
– 消费级显卡（如 RTX 系列）可能需要单独处理 INT8 量化

避免频繁重新编译引擎：设置合理的 profile 范围

警惕 shape 变化导致的并行度下降：

# 错误示例 - 不同 batch 的 shape 变化过大
profile.set_shape("input", (1,3,224,224), (32,3,512,512), (64,3,1024,1024))

# 正确做法 - 分阶段处理不同分辨率
if img_size < 256:
    engine = small_engine
else:
    engine = large_engine