PyTorch Profiler 性能优化示例：定位 TorchMetrics 收集瓶颈，提高 GPU 利用率

指标收集是每个机器学习项目不可或缺的组成部分，它使我们能够跟踪模型性能并监控训练进度。理想情况下，我们希望在不给训练过程带来额外开销的前提下收集和计算指标。与训练循环的其他部分一样，低效的指标计算可能会引入不必要的开销，延长训练步骤的耗时，并增加训练成本。

本文是将聚焦于指标收集，演示指标收集的一种简单实现如何对运行时性能产生负面影响，并探讨用于分析和优化它的工具与技术。

为了实现指标收集，我们将使用 TorchMetrics，这是一个标准化 PyTorch 中指标计算的常用库。我们的目标是：

为了便于讨论，我们将定义一个简单的 PyTorch 模型，并评估指标收集对运行时性能的影响。我们将在 NVIDIA A40 GPU 上运行实验，使用的 PyTorch 版本为 2.5.1（通过 Docker 镜像），TorchMetrics 版本为 1.6.1。

需要注意的是：指标收集的行为可能因硬件、运行时环境和模型架构而异。本文中提供的代码片段仅用于演示目的。

 

https://avoid.overfit.cn/post/429185b1fd4f49eea7d4140174cce7a2