13、深度学习性能调优：火焰图与GPU内存调试

深度学习性能调优：火焰图与GPU内存调试

1. 识别CPU瓶颈与火焰图的使用

在实际的深度学习解决方案中，CPU瓶颈经常出现。例如，许多基于图像的变换操作使用Python Imaging Library，这些操作完全受CPU限制。在处理大型数据集时，训练循环中会反复进行这些变换，这可能成为性能瓶颈。火焰图是解决此类问题的有力工具，它能帮助我们直观地了解程序的执行情况。

1.1 安装py - spy

生成可转换为火焰图的堆栈跟踪有多种方法，这里我们选择使用基于Rust的堆栈分析器py - spy，它可以直接生成火焰图。安装命令如下：

pip install py-spy

你可以通过 --pid 参数找到正在运行的进程的进程标识符（PID）并附加py - spy：

py-spy --flame profile.svg --pid 12345

或者传入一个Python脚本，以下是一个简单的Python脚本示例：

import torch
import torchvision

def get_model():
    return torchvision.models.resnet18(pretrained=True)

def get_pred(model):
    return model(torch.rand([1,3,224,224]