PyTorch安装

发表于 分类于 阅读次数： 评论数： 本文字数： 515 阅读时长 ≈ 2 分钟
PyTorch安装

1. 检查 NVIDIA GPU 驱动是否已安装

检查 GPU 驱动版本

打开终端或命令行，运行以下命令：

nvidia-smi

输出示例：

Thu Nov 28 11:03:22 2024
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 560.81                 Driver Version: 560.81         CUDA Version: 12.6     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                  Driver-Model | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA GeForce RTX 3060 ...  WDDM  |   00000000:01:00.0 Off |                  N/A |
| N/A   54C    P0             24W /   71W |       0MiB /   6144MiB |      0%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+

+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI        PID   Type   Process name                              GPU Memory |
|        ID   ID                                                               Usage      |
|=========================================================================================|
|  No running processes found                                                             |
+-----------------------------------------------------------------------------------------+
  • 如果能看到类似输出，说明 GPU 驱动已安装。
  • 如果命令不存在或未安装驱动，可以前往 NVIDIA 官网下载驱动，选择您的 GPU 型号下载并安装。

2. 安装支持 GPU 的 CUDA 工具链

检查 CUDA 版本

  1. 确保 CUDA 已安装：

    nvcc --version

    输出示例：

    nvcc: NVIDIA (R) Cuda compiler driver
    Copyright (c) 2005-2024 NVIDIA Corporation
    Built on Thu_Mar_28_02:30:10_Pacific_Daylight_Time_2024
    Cuda compilation tools, release 12.4, V12.4.131
    Build cuda_12.4.r12.4/compiler.34097967_0

  2. 如果未安装 CUDA：

    • 前往 CUDA 下载页面。
    • 根据 GPU 支持的 CUDA 版本和系统配置下载对应的 CUDA 工具包。

3. 安装支持 GPU 的 PyTorch

PyTorch 的 CUDA 支持与您的 CUDA 工具链版本相关，需要选择匹配的版本。

安装最新 PyTorch 版本

  1. 前往 PyTorch 官网。
  2. 根据您的配置生成安装命令。

示例安装命令：

如果您使用的是 Python 环境，CUDA 版本为 12.4，运行以下命令：

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

验证 PyTorch 是否正确安装并支持 GPU：

  1. 运行以下代码：

    import torch
    print("PyTorch 版本:", torch.__version__)
    print("CUDA 是否可用:", torch.cuda.is_available())
    print("GPU 名称:", torch.cuda.get_device_name(0) if torch.cuda.is_available() else "未检测到 GPU")

  2. 输出示例：

    PyTorch 版本: 2.5.1+cu124
    CUDA 是否可用: True
    GPU 名称: NVIDIA GeForce RTX 3060 Laptop GPU

4. 测试完整 GPU 环境

运行以下代码验证环境：

import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"当前设备: {device}")
x = torch.rand(3, 3).to(device)
print("在 GPU 上的张量:", x)

如果没有报错并成功输出张量，说明 GPU 环境已配置正确。