Python Pytorch 环境配置

记录一下当前 Python 涉及的一些环境配置。

基本信息

当前本地以及服务器的Python和CUDA的基本信息如下：

• Windows（以及WSL）：
  • Windows 11
  • miniconda
  • Python 版本为 3.12.1
  • CUDA 驱动最高支持版本 12.8（nvidia-smi）
• GPU服务器：
  • Ubuntu 20.04
  • anaconda
  • Python 版本为 3.11.7
  • CUDA 驱动最高支持版本 12.4

两者存在很多差异，为了尽量保持本地和服务器的版本一致，选择使用 Python 3.12.x，CUDA 12.4。

说明：

• 指定 Python 版本是 conda 相比于 venv 的一大优势，另一个优势是 conda 包不局限于 Python。
• 很多库并不支持这么高版本的 Python，到时候再说，反正版本问题是避免不了的。

常用包

记录一些常用的包

# 必备
conda install numpy scipy pandas matplotlib seaborn scikit-learn sympy jupyter

# 支持import ipynb文件
pip install import-ipynb

# 进度条
conda install tqdm

PyTorch 安装

目前 PyTorch 官网 的稳定版为 2.6.0，要求：

• python >= 3.9
• 支持 cuda 12.4、cuda 12.6 或 cpu 版本
• 建议使用 pip 安装，只提供 pip 安装命令，不再提供 conda 安装命令（可以在切换到 conda 环境之后，使用对应的 pip 安装）

官网提供的几种安装命令如下

# (*) windows cuda=12.4
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124

# windows cuda=12.6
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126

# windows cpu
pip3 install torch torchvision torchaudio

# (*) linux cuda=12.4
pip3 install torch torchvision torchaudio

# linux cuda=12.6
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126

# linux cpu
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu

解释一下：

• torch: 是 PyTorch 基础库，提供张量计算和自动求导等基础功能。
• torchvision: 是 PyTorch 的一个子模块，专门用于视觉处理任务。
• torchaudio: 是 PyTorch 的另一个子模块，专门用于音频处理任务。

其实还有一个已经停止更新的 torchtext。

conda 环境配置记录

Windows 的 conda 环境 myenv-pytorch 的配置记录：

# 创建环境
conda create --name myenv-pytorch python=3.12.1

# 激活环境
conda activate myenv-pytorch

# 确保在激活的conda环境中使用pip
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124

# 安装常用包
conda install numpy scipy pandas matplotlib seaborn scikit-learn sympy jupyter

GPU服务器的 conda 环境 myenv-pytorch-linux 的配置记录：

# 创建环境
conda create --name myenv-pytorch-linux python=3.12.7

# 激活环境
conda activate myenv-pytorch-linux

# 确保在激活的conda环境中使用pip
pip3 install torch torchvision torchaudio

# 安装常用包
conda install numpy scipy pandas matplotlib seaborn scikit-learn sympy jupyter

注意：

• 不要随便升级numpy，这可能导致pytorch无法正常运行。
• 上面两个conda环境中实际选择的numpy版本不一样，可能是python版本和平台差异导致的。

PyTorch CUDA 测试脚本

可以跑一段 Python 脚本来检测当前环境安装的 PyTorch 是否正常运行，是否支持使用CUDA，以及检测GPU的信息

import torch

def check_gpu():
    if torch.cuda.is_available():
        num_gpus = torch.cuda.device_count()
        print(f"CUDA is available! Number of GPUs: {num_gpus}\n")

        for i in range(num_gpus):
            prop = torch.cuda.get_device_properties(i)
            print(f"GPU {i} Name: {prop.name}")
            print(f"GPU {i} Total Memory: {prop.total_memory / (1024 ** 3):.2f} GB")
            print(f"GPU {i} Compute Capability: {prop.major}.{prop.minor}\n")
    else:
        print("CUDA is not available.")

if __name__ == "__main__":
    check_gpu()

例如在个人笔记本的输出

CUDA is available! Number of GPUs: 1

GPU 0 Name: NVIDIA GeForce RTX 4060 Laptop GPU
GPU 0 Total Memory: 8.00 GB
GPU 0 Compute Capability: 8.9

在GPU服务器上的输出

CUDA is available! Number of GPUs: 6

GPU 0 Name: NVIDIA RTX A6000
GPU 0 Total Memory: 47.53 GB
GPU 0 Compute Capability: 8.6

GPU 1 Name: NVIDIA RTX A6000
GPU 1 Total Memory: 47.53 GB
GPU 1 Compute Capability: 8.6

GPU 2 Name: NVIDIA RTX A6000
GPU 2 Total Memory: 47.53 GB
GPU 2 Compute Capability: 8.6

GPU 3 Name: NVIDIA RTX A6000
GPU 3 Total Memory: 47.53 GB
GPU 3 Compute Capability: 8.6

GPU 4 Name: NVIDIA RTX A6000
GPU 4 Total Memory: 47.53 GB
GPU 4 Compute Capability: 8.6

GPU 5 Name: NVIDIA RTX A6000
GPU 5 Total Memory: 47.53 GB
GPU 5 Compute Capability: 8.6

PyTorch 指定 GPU

为了指定 Pytorch 使用的 GPU，在代码中通常有如下语句

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

如果 GPU 可用，默认使用编号为 0 的物理 GPU，在没有 GPU 的情况下自动退化为使用 CPU，

对于多个GPU的情形，可以手动加上编号来指定使用的GPU，例如cuda:0，cuda:1，cuda:2等，例如

device = 'cuda:1'

注意：虽然cuda:0和cuda都指代第一个显卡，但是直接进行比较仍然可能被视作不同设备。

可以在调用命令之前设置环境变量CUDA_VISIBLE_DEVICES，它可以直接控制 PyTorch / TensorFlow / CUDA 程序能看见哪几张 GPU，但是作用范围仅限于当前命令（除非export导出环境变量）。

例如限制使用单个 GPU

CUDA_VISIBLE_DEVICES=1 python train.py

此时程序只能看见和使用这张卡：程序中的 torch.cuda.device_count() = 1，cuda:0 对应编号为 1 的物理 GPU。

限制使用多个 GPU

CUDA_VISIBLE_DEVICES=2,3 python train.py

此时程序只能看见和使用这两张卡：程序中的 torch.cuda.device_count() = 2，cuda:0 对应编号为 2 的物理 GPU，cuda:1 对应编号为 3 的物理 GPU。

可以使用下面的工具函数来使用：

• 指定编号的GPU；
• 目前最空闲的GPU（原理是在shell中调用 nvidia-smi 获取状态信息）；
• 所有可用的GPU。

import subprocess
import torch


def try_gpu(i=0):
    if torch.cuda.device_count() >= i + 1:
        return torch.device(f'cuda:{i}')
    return torch.device('cpu')

def try_all_gpus():
    devices = [torch.device(f'cuda:{i}')
             for i in range(torch.cuda.device_count())]
    return devices if devices else [torch.device('cpu')]

def try_free_gpu():
    if not torch.cuda.is_available():
        return torch.device("cpu")

    try:
        result = subprocess.check_output(
            ["nvidia-smi", "--query-gpu=memory.used", "--format=csv,nounits,noheader"],
            encoding="utf-8",
        )
        memory_usage = [int(x) for x in result.strip().split("\n")]

        best_gpu = min(range(len(memory_usage)), key=lambda i: memory_usage[i])
        return torch.device(f"cuda:{best_gpu}")

    except Exception:
        print("Failed to query GPU info, fallback to cuda:0")
        return torch.device("cuda:0")

device = try_gpu(1)
device = try_free_gpu()
devices = try_all_gpus()

至于并行使用多个 GPU 的情况，暂时不作讨论。

在深度学习模型中，对于使用固定尺寸的卷积层，cuDNN 提供了多个卷积算法实现，可以开启下面的选项来评估并选择最优的算法，选项默认关闭

if torch.cuda.is_available():
    torch.backends.cudnn.benchmark = True