深度学习训练与推理的性能分析（Profiling）

在深度学习的工程实践中，性能瓶颈往往不是一眼就能看出来的。无论是训练还是推理阶段，都需要借助 profiling 工具来定量分析各个环节的耗时，才能有针对性地做优化。这篇笔记记录了我在训练和推理两个阶段分别用到的 profiling 方法。

训练过程的 Profiling

训练阶段的 profiling 主要目的是分析数据加载、前向传播、反向传播各部分的时间占用，找出整个训练流程里的时间瓶颈。

在 PyTorch 中，我用的是 NVTX（NVIDIA Tools Extension）来做标注。在训练循环里，针对每个主要阶段（数据加载、前向、loss 计算、反向、optimizer 更新）前后插入标记，代码写成如下这样：

import torch.cuda.nvtx as nvtx

def train(args, model, device, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        nvtx.range_push("Batch " + str(batch_idx))

        nvtx.range_push("Copy to device")
        data, target = data.to(device), target.to(device)
        nvtx.range_pop()

        nvtx.range_push("Forward pass")
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        nvtx.range_pop()
        nvtx.range_pop()

经过标注之后，用 Nsight Systems 打开采集结果，时间线会按 NVTX 区间着色，各阶段耗时一目了然。自己测试了一下，实际效果如下图所示，确实挺有用的，再也不用靠打印时间戳来做粗粒度估计了：

推理过程的 Profiling

推理阶段的情况稍有不同。通常的做法是先把 PyTorch 训练好的模型转换为 TensorRT，再用 TensorRT 进行推理部署。针对这个流程的 profiling，需要在 TensorRT 的上下文中进行，与训练阶段的 NVTX 标注方式有所区别，后续打算进一步整理这部分的具体用法。