如何在Pytorch中精细化利用显存

### 关于如何使用 PyTorch 配合 NVIDIA Nsight Systems NVIDIA Nsight Systems 是一种性能分析工具，能够帮助开发者深入了解应用程序的行为以及优化其运行效率。当与 PyTorch 结合时，它可以帮助识别瓶颈并改进模型训练过程中的资源利用率。 #### 安装环境准备 为了确保可以顺利使用 PyTorch 和 Nsight Systems 进行联合调试，需满足以下条件： - 确认安装的 CUDA 版本支持所选的 PyTorch 版本[^1]。 - 下载适合目标平台的 PyTorch 轮子文件（whl），例如 `torch-1.8.0-cp36-cp36m-linux_aarch64.whl` 并完成安装。 ```bash wget https://nvidia.box.com/shared/static/p57jwntv436lfrd78inwl7iml6p13fzh.whl -O torch-1.8.0-cp36-cp36m-linux_aarch64.whl sudo apt-get install python3-pip libopenblas-base libopenmpi-dev pip3 install Cython pip3 install numpy torch-1.8.0-cp36-cp36m-linux_aarch64.whl ``` #### 启动 Nsight Systems 对 PyTorch 应用程序进行剖析 Nsight Systems 提供了一个图形界面来捕获和可视化应用的工作负载数据。以下是具体操作方法： 1. **启动 Nsight Systems 工具** 打开终端输入命令启动 GUI 或者通过桌面快捷方式打开。 2. **配置捕捉选项** 在 Nsight Systems 中设置要监控的目标进程及其参数。对于 Python 脚本来说，通常指定完整的路径到脚本文件即可。 3. **执行捕捉动作** 开始记录会话之后，让您的 PyTorch 训练循环正常运行一段时间以便收集足够的样本数据。 4. **审查报告结果** 停止捕捉后，软件将生成一份详细的交互图表展示 CPU/GPU 协作状况以及其他重要指标如显存带宽消耗等信息[^3]。 #### 利用其他辅助工具增强诊断能力 除了直接依赖 Nsight Systems 外部功能外，还可以结合一些内置 API 来获取更精细的信息: - 使用 `torch.cuda.memory_allocated()` 查询当前分配了多少字节给张量对象存储空间[^2]； - 尝试调用 `torch.utils.benchmark.Timer` 测定特定代码片段耗时时长； - 如果涉及复杂网络结构定义，则推荐尝试第三方库比如 `torchinfo` 自动生成逐层统计摘要表单； ```python import torch print(f"Allocated Memory: {torch.cuda.memory_allocated() / 1e9:.2f} GB") from torch.utils.benchmark import Timer timer = Timer(stmt='model(input)', globals={'model': model, 'input': input}) res = timer.timeit(100) print(res) ``` ---