Skip to content

PyTorch

Pytorch 官网 https://pytorch.org/

安装Pytorch

在安装pytorch 之前请创建一个新的环境 Conda 创建虚拟环境

  1. 请根据自己的cuda版本选择安装 torch版本
nvcc -V

# 输出类似于
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2022 NVIDIA Corporation
Built on Wed_Sep_21_10:41:10_Pacific_Daylight_Time_2022
Cuda compilation tools, release 11.8, V11.8.89
Build cuda_11.8.r11.8/compiler.31833905_0

如果没有输出,可以查看如何安装Linux 配置CUDA 环境变量

  1. cuda 版本 为11.8 在这儿的pytorch版本上我也选择对应版本

  • 复制下面生产的命令运行即可
    pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
    # 或 建议使用pip,# pip 安装默认使用系统的代理方式
    #conda 安装往往时间比较长,而且设置代理可能不会太方便
    # conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia
    
    如果在设置 conda 国内源下,下载仍慢可以设置Conda 代理

其他

脚本调试指定CUDA设备

首先导入os ,然后定义对应的环境变量

import os  

if 'CUDA_VISIBLE_DEVICES' not in os.environ: # 兼容如果使用sh 脚本设置多卡被覆盖
    # 设置使用的GPU,
    os.environ['CUDA_VISIBLE_DEVICES'] = '2'

# 其他代码或设置其他环境变量....