如何设置 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 代码。

在 Linux 或 macOS 上设置环境变量

在终端中，你可以使用 export 命令来设置环境变量：

export TORCH_CUDA_ARCH_LIST="8.9"

在 Windows 上设置环境变量

在命令提示符（CMD）中，你可以使用 set 命令：cmd

set TORCH_CUDA_ARCH_LIST=8.9

在 PowerShell 中，你可以使用：

$env:TORCH_CUDA_ARCH_LIST="8.9"

构建优化的 PyTorch 版本

设置好环境变量后，你就可以开始构建针对特定 CUDA 架构优化的 PyTorch 版本了。这对于确保你的深度学习模型能够充分利用 GPU 的性能至关重要。

结论

正确设置 TORCH_CUDA_ARCH_LIST 环境变量可以显著提高你的 PyTorch 应用的性能。通过遵循上述步骤，你可以确保你的深度学习模型在 GPU 上运行得更快、更高效。如果你在设置过程中遇到任何问题，不要犹豫，查阅 PyTorch 官方文档或寻求社区的帮助。