如何设置 TORCH_CUDA_ARCH_LIST 环境变量以优化 PyTorch 性能

引言

在深度学习领域，PyTorch 是一个广泛使用的框架，它允许开发者高效地构建和训练模型。为了充分利用你的 GPU 硬件，正确设置 TORCH_CUDA_ARCH_LIST 环境变量至关重要。这个变量告诉 PyTorch 在构建过程中应该针对哪些 CUDA 架构版本进行优化。本文将指导你如何确定你的 GPU 的 CUDA 架构能力，并设置相应的环境变量。

确定你的 GPU 的 CUDA 架构能力

首先，你需要知道你的 GPU 支持的 CUDA 计算能力。你可以通过运行以下 Python 代码来获取这个信息：

import torch; print(torch.cuda.get_device_capability())

或者，如果你更喜欢使用命令行，可以执行：

python -c "import torch; print(torch.cuda.get_device_capability())"

这将返回一个元组，包含两个整数，分别代表你的 GPU 支持的 CUDA 架构的主版本号和次版本号。例如，如果输出是 (8, 9)，则表示你的 GPU 支持 CUDA 架构 8.9。

设置 TORCH_CUDA_ARCH_LIST 环境变量

一旦你知道了你的 GPU 的 CUDA 架构能力，你就可以设置 TORCH_CUDA_ARCH_LIST 环境变量，以便 PyTorch 可以针对这些架构进行优化。这个列表告诉 PyTorch 你的 GPU 支持的 CUDA 版本，以便正确编译和优化 PyTorch 代码。

在 Linu