ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

PyTorch 1.8与CUDA 11.1环境配置实战:从安装到排错

PyTorch 1.8与CUDA 11.1环境配置实战:从安装到排错 1. 版本兼容性是第一道生死线1.1 为什么选 CUDA 11.1 PyTorch 1.8 这个组合先说结论这个组合在 2021 年到 2023 年之间几乎是深度学习炼丹师的“标配”之一。PyTorch 1.8 是 2021 年 3 月发布的版本官方在发布时明确支持 CUDA 11.1 和 CUDA 10.2 两套运行时。CUDA 11.1 是 NVIDIA 在 2020 年底推出的版本当时对 Ampere 架构也就是 RTX 30 系列显卡做了完整适配像 A100、RTX 3090 这些卡在 CUDA 11.1 下跑得最稳。放到今天来看虽然 CUDA 12.x 和 PyTorch 2.x 早就出来了但很多老项目、论文复现、课程作业还锁死在 PyTorch 1.8 上。尤其是一些工业界的推理服务为了“稳定压倒一切”宁可守着旧版本也不升级。所以这个组合并没有过时反而因为踩坑的人足够多解决方案也变得非常成熟。但我必须说句实在话在动手配置之前先想清楚你到底需不需要这个组合。如果你的显卡是 RTX 30 系列之前的老卡比如 GTX 10 系列CUDA 11.1 反而支持得没毛病如果是 RTX 40 系列比如 4090那 CUDA 11.1 根本装不上驱动适配层这时候你需要的不是这套环境而是先升级到 CUDA 12.x 或者把 PyTorch 升到 1.13 以上。判断标准很简单去 NVIDIA 官网查你显卡的Compute Capability算力代号如果大于等于 8.0那 CUDA 11.1 就是可用的如果是 8.6 以上建议直接上 CUDA 11.8。1.2 最容易忽略的版本对应关系很多人一上来就 pip install torch1.8.0结果装完发现 torch.cuda.is_available() 返回 False直接心态崩了。这里有个大坑PyTorch 的 PyPI 默认包是带 CUDA 10.2 运行时的你要的是 CUDA 11.1 版本必须从官方指定的源安装。PyTorch 官方给出来的安装命令是这样的pip install torch1.8.0cu111 torchvision0.9.0cu111 torchaudio0.8.0 -f https://download.pytorch.org/whl/torch_stable.html注意看这个命令里的 cu111 后缀这才是关键。如果你直接 pip install torch1.8.0默认拿到的是 CPU 版本或者 CUDA 10.2 版本跟你系统里装的 CUDA 11.1 根本对不上。除了 PyTorch 本身还要注意两个配套库的版本库名对应版本说明torchvision0.9.0cu111必须和 torch 一一对应错一个版本都可能引发 C 库不匹配torchaudio0.8.0音频处理库不做音频可以暂时不装但装了就得对版python3.6 - 3.8PyTorch 1.8 时代对 Python 3.9 的适配还很不完善这里特别提醒一句Python 版本别用最新的。PyTorch 1.8 发布的时候 Python 3.9 刚出来不久官方轮子对 3.9 的支持处于“半残”状态。我见过太多人用 Python 3.9 装 torch 1.8装是能装上一跑就报 OSError 或者段错误。最稳的选择是 Python 3.8.10退一步用 3.7 也没问题。2. 环境准备显卡驱动、CUDA Toolkit、cuDNN 三者到底什么关系2.1 驱动、Toolkit、cuDNN 的角色分工很多新手把“CUDA 版本”理解成“显卡驱动的版本”这是个大误区。我把这三个东西的关系用大白话说清楚显卡驱动是操作系统和显卡硬件之间的翻译官。它决定了你的显卡能不能被系统识别以及能跑多新的 CUDA 运行时。CUDA Toolkit是一整套开发工具包包含编译器nvcc、各种库文件如 cublas、cudart和头文件。你系统里装了这个才能编译 CUDA 代码也才能让 PyTorch 找到 CUDA 运行时。cuDNN是 NVIDIA 专门为深度学习优化的深度神经网络库卷积运算、LSTM 等都靠它加速。PyTorch 不会直接调用你系统里的 cuDNNconda 安装时会单独带一份但如果你要自己编译 C 扩展或者用 TensorRT那系统里的 cuDNN 版本就得对上。简单来说驱动是地基CUDA Toolkit 是脚手架cuDNN 是装修材料。PyTorch 在 conda 模式下会自带一份 CUDA 运行时和 cuDNN不需要你系统里装了完整的 CUDA Toolkit 也能跑但在 pip 模式下PyTorch 虽然也内置了部分运行时编译自定义算子时还是得依赖系统的 CUDA。2.2 显卡驱动的安装与检查如果你已经装了 NVIDIA 驱动先别急着装 CUDA。打开命令行输入nvidia-smi看右上角那行“CUDA Version”这里显示的是当前驱动所支持的最高 CUDA 版本。比如显示 11.4那就意味着你的驱动可以兼容 CUDA 11.1 及以下的所有 CUDA 版本。注意这里的 CUDA Version 不是你已经装好的 CUDA 版本而是驱动的能力上限。如果这个数值小于 11.1那你需要先升级显卡驱动。升级驱动的方式到 NVIDIA 官网驱动下载页面输入你的显卡型号和操作系统。下载 Game Ready 或 Studio 驱动都行推荐 Studio 驱动稳定性更好。安装时选择“自定义安装”勾选“执行清洁安装”避免旧驱动残留。这里有个坑很多笔记本用户有“双显卡”Intel 集显 NVIDIA 独显装驱动的时候系统可能默认给集显装驱动导致 nvidia-smi 根本找不到命令。解决办法是去设备管理器里确认 NVIDIA 显卡是否存在如果存在但 nvidia-smi 不可用就先把集显驱动禁掉装完 NVIDIA 驱动再启用。2.3 CUDA Toolkit 11.1 安装的两种模式CUDA Toolkit 11.1 的安装方式主要有两类本地 runfileLinux和 exeWindows以及网络版安装器。我强烈推荐本地安装包网络版安装器在网速不稳的情况下很容易下载到一半就失败断了还得重来。Linux 下的 runfile 安装步骤wget https://developer.download.nvidia.com/compute/cuda/11.1.0/local_installers/cuda_11.1.0_455.23.05_linux.run sudo sh cuda_11.1.0_455.23.05_linux.run这里一定要注意安装过程中会让你选择是否安装显卡驱动。如果你的驱动已经装好了且版本大于等于 455.23.05务必在安装界面取消勾选 Driver只选 CUDA Toolkit 那一项否则会把当前驱动给替换掉搞不好就黑屏了。Windows 下比较简单双击 exe 安装包选择“自定义”安装模式。同样如果驱动已有取消勾选“Display Driver”组件保留 CUDA 相关的所有组件。安装路径建议选默认的C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.1不要乱改路径因为后续很多工具会基于默认路径去查找。安装完成后检查是否成功nvcc --version如果输出里有 “Cuda compilation tools, release 11.1” 字样说明 CUDA Toolkit 装好了。2.4 环境变量配置——容易出问题的地方装完 CUDA Toolkit 后环境变量不配置会导致一个经典报错nvcc -V找不到命令或者 PyTorch 编译扩展时提示找不到cuda.h。Linux 下在~/.bashrc或~/.zshrc末尾追加export PATH/usr/local/cuda-11.1/bin:$PATH export LD_LIBRARY_PATH/usr/local/cuda-11.1/lib64:$LD_LIBRARY_PATH然后执行source ~/.bashrc让配置生效。Windows 下在“系统属性 → 环境变量”中在Path变量里新增C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.1\bin新增系统变量CUDA_HOME值为C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.1如果使用 Visual Studio 编译 CUDA 代码还需要在INCLUDE和LIB环境变量里加上 CUDA 的 include 和 lib 目录不过这只是针对 C 开发PyTorch 用户一般不用手动配置。这里有个很多人会忽略的细节如果你系统里之前装过 CUDA 10.2 或者其他版本/usr/local/cuda这个软链接指向的未必是 11.1。用ls -l /usr/local/cuda看一眼如果指向的不是 v11.1可以手动改软链接sudo rm -rf /usr/local/cuda sudo ln -s /usr/local/cuda-11.1 /usr/local/cuda不推荐这样直接删软链接的做法更稳妥的是用update-alternatives管理多版本但为了简单省事删软链接重新指是大多数人的做法前提是你只保留 CUDA 11.1 这一个版本。2.5 cuDNN 的安装与验证cuDNN 需要去 NVIDIA 官网注册账号才能下载这一步经常卡住很多人。下载时选择“cuDNN for CUDA 11.x”对应的版本推荐 8.0.5 或 8.1.x 都可以。安装本质上就是解压复制文件Linux 下执行tar -xzvf cudnn-11.3-linux-x64-v8.1.0.77.tgz sudo cp cuda/include/cudnn.h /usr/local/cuda-11.1/include/ sudo cp cuda/lib64/libcudnn* /usr/local/cuda-11.1/lib64/ sudo chmod ar /usr/local/cuda-11.1/include/cudnn.h /usr/local/cuda-11.1/lib64/libcudnn*Windows 下就是把解压出来的 bin、include、lib 目录里的文件分别复制到 CUDA 安装目录的对应文件夹中。注意 Windows 下执行复制文件操作时建议用管理员权限的终端否则可能提示权限不足。验证 cuDNN 是否生效cat /usr/local/cuda-11.1/include/cudnn.h | grep CUDNN_MAJOR -A 2如果输出里有版本号信息说明安装成功。还有个小细节如果你用的是 conda 安装 PyTorch系统里的 cuDNN 版本是“锦上添花”而不是“雪中送炭”。因为 conda 版 PyTorch 自带 cudatoolkit 和 cudnn你系统里装不装 cuDNN 它都能跑。但如果你后面要用 NVIDIA TensorRT、DeepStream 这类工具系统 cuDNN 就必须存在。3. PyTorch 1.8 安装实操conda 与 pip 路线之争3.1 conda 安装省心但慢进入正题。PyTorch 1.8 CUDA 11.1 的安装有两条路线。先看 conda 路线conda create -n torch18 python3.8 -y conda activate torch18 conda install pytorch1.8.0 torchvision0.9.0 torchaudio0.8.0 cudatoolkit11.1 -c pytorch -c conda-forge这条命令会把 CUDA 11.1 的运行时cudatoolkit和配套的 cuDNN 一起装到你的 conda 环境里不需要提前在系统里装 CUDA Toolkit。对于只想跑模型、不打算编译 CUDA 扩展的人来说这是最省事的路线。但 conda 路线最大的痛点是下载速度。PyTorch 的 conda 包比较大500MB 左右国内下载经常卡在进度条不动。解决方法是配置国内 conda 镜像比如清华源conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/pytorch/ conda config --set show_channel_urls yes不过有一点得提醒清华源同步 PyTorch 官方 conda 频道不是实时的某些特定版本可能滞后。如果配了镜像还是下载失败就临时把 channel 指回官方源conda install pytorch1.8.0 torchvision0.9.0 torchaudio0.8.0 cudatoolkit11.1 -c pytorch -c conda-forge3.2 pip 安装可控但容易踩版本坑pip 路线的安装命令pip install torch1.8.0cu111 torchvision0.9.0cu111 torchaudio0.8.0 -f https://download.pytorch.org/whl/torch_stable.html这条命令会从 PyTorch 官方轮子仓库下载带 CUDA 11.1 运行时编译的包。和 conda 的区别在于pip 版 PyTorch 内置了 CUDA 运行时但并没有把整个 CUDA Toolkit 都打包进去。如果你想在 PyTorch 里用torch.utils.cpp_extension编译自定义 CUDA 算子还是得自己装完整的 CUDA Toolkit。pip 路线的优点是版本控制更精细通过 cu111 后缀可以精确锁定你要的 CUDA 版本不受 conda 频道同步延迟的影响。缺点是对系统依赖要求更高比如 GCC 版本不能太高建议 GCC 7 或 8否则编译扩展时会出现各种奇怪的宏定义错误。国内用户下载 pip 包太慢的问题可以用国内 PyPI 镜像pip install torch1.8.0cu111 torchvision0.9.0cu111 torchaudio0.8.0 -f https://download.pytorch.org/whl/torch_stable.html -i https://pypi.tuna.tsinghua.edu.cn/simple注意-i指定的是 PyPI 镜像地址-f指定的是 PyTorch 轮子仓库地址两者不冲突。但有一个坑某些镜像站可能没缓存 PyTorch 的 cu111 后缀轮子所以最稳妥的做法是先用-i下载依赖小包torch 本体还是走官方源。3.3 安装完的第一时间验证不管用哪条路线装完之后都要第一时间验证环境是否真的可用。创建一个 Python 脚本test_torch.pyimport torch print(PyTorch version:, torch.__version__) print(CUDA available:, torch.cuda.is_available()) if torch.cuda.is_available(): print(CUDA version:, torch.version.cuda) print(cuDNN version:, torch.backends.cudnn.version()) print(GPU count:, torch.cuda.device_count()) print(GPU name:, torch.cuda.get_device_name(0))在命令行执行python test_torch.py理想输出长这样PyTorch version: 1.8.0cu111 CUDA available: True CUDA version: 11.1 cuDNN version: 8005 GPU count: 1 GPU name: NVIDIA GeForce RTX 3090如果你的输出里CUDA available是 False那就进入下一章这是大部分人真正卡住的地方。4. 高频报错与排查方法实战4.1 torch.cuda.is_available() 返回 False这是最常见的报错没有之一。出现这个问题的原因有几类我按出现频率从高到低排列第一类PyTorch 装成了 CPU 版。输入pip list | grep torch看一下 torch 的版本信息。如果显示的是torch1.8.0没有 cu111 后缀说明你装的是 CPU 版。卸载重装pip uninstall torch torchvision torchaudio -y pip install torch1.8.0cu111 torchvision0.9.0cu111 torchaudio0.8.0 -f https://download.pytorch.org/whl/torch_stable.html第二类显卡驱动版本太低。用nvidia-smi检查驱动支持的 CUDA 版本如果低于 11.1PyTorch 无法通过驱动访问显卡。升级驱动到 455.23 以上版本推荐 470 或 510 系列更稳。第三类驱动没有正确安装。这种情况常见于笔记本双显卡用户NVIDIA 驱动装上了但系统默认走集显渲染。用nvidia-smi如果报错说明驱动真的没装好需要重新安装。第四类conda 环境混用。你在 conda 环境 A 里安装了 PyTorch但激活的确是环境 B。或者你之前在环境 A 里用 pip 装过包又用 conda 装了冲突的包。解决方案是新建一个干净环境重装。4.2 安装时提示 .run 文件格式错误或校验失败前文提到下载 CUDA Toolkit 时很多人会遇到这个报错gzip: stdin: invalid compressed># 先看文件大小是否和官网一致 ls -lh cuda_11.1.0_455.23.05_linux.run # 用 md5sum 校验 md5sum cuda_11.1.0_455.23.05_linux.run和官网给出的 MD5 值对比不一致就重新下载。如果 wget 想断点续传执行wget -c https://developer.download.nvidia.com/compute/cuda/11.1.0/local_installers/cuda_11.1.0_455.23.05_linux.run4.3 运行时报错libcudart.so.11.0: cannot open shared object file明明装好了 CUDA跑 PyTorch 却报找不到 libcudart.so.11.0这就是环境变量没配好的问题。检查/usr/local/cuda-11.1/lib64 目录下有没有 libcudart.so.11.0 这个文件注意版本号是 .so.11.0不是 .so.10.2。LD_LIBRARY_PATH 是否包含这个目录。有时候环境变量配了但没生效需要重启终端或者重新source ~/.bashrc。如果用的是 Anaconda还可以用 conda 包直接补全conda install cudatoolkit11.1 -c conda-forge这样会在 conda 环境里放一份 CUDA 运行时PyTorch 在 conda 环境下会自动优先查找环境内的库。4.4 安装 PyTorch 时提示 Cannot find suitable CUDA version有些人在 pip 安装时看到 “Cannot find suitable CUDA version, but this package is not CUDA-enabled” 的警告然后装出来的 PyTorch 就是 CPU 版。这种情况一般发生在你使用旧版 pip 时pip 无法解析cu111这种非标准版本号。升级 pip 即可解决pip install --upgrade pip然后重新执行安装命令。如果你的 pip 升级到了 21 以上版本一般就不会再出现这个警告了。4.5 多 CUDA 版本切换与冲突很多人系统里不只有一个 CUDA 版本比如后来为了新项目装了 CUDA 11.8结果 PyTorch 1.8 的环境突然跑不了了。原因很简单环境变量 PATH 和 LD_LIBRARY_PATH 指向了版本 A但 PyTorch 需要版本 B。解决方案是做好版本隔离。既然你已经用 conda 管理 Python 环境那就不要让系统全局的 CUDA 干扰到 conda 环境。具体操作# 在激活 conda 环境后强制指定 CUDA 相关路径为环境内版本 conda activate torch18 export PATH/usr/local/cuda-11.1/bin:$PATH export LD_LIBRARY_PATH/usr/local/cuda-11.1/lib64:$LD_LIBRARY_PATH更推荐的做法是完全依赖 conda 环境内的 cudatoolkit系统层面的 CUDA 版本只服务于系统级的编译需求两者互不干扰conda install cudatoolkit11.1 -c conda-forge这样 PyTorch 在 conda 环境下会自动找到环境内的 cudatoolkit不会去碰系统的 /usr/local/cuda。4.6 编译自定义 CUDA 扩展时报错如果你要在 PyTorch 1.8 里编译自定义算子比如用 torch.utils.cpp_extension.load 加载 .cu 文件报错往往千奇百怪最常见的是error: identifier __CUDA_ARCH__ is undefined或者/usr/include/c/8/bits/c17_warning.h:32:2: error: #error This file requires compiler and library support for the ISO C 2011这些问题的根源在于 GCC 和 CUDA 版本的兼容性。CUDA 11.1 官方的 nvcc 最高支持 GCC 9.x但 PyTorch 1.8 的编译环境对 GCC 版本很敏感。如果你的系统 GCC 是 10 或 11 版本编译扩展时就会出问题。解决方案在编译命令前显式指定低版本 GCC如果系统装了旧版 GCCexport CC/usr/bin/gcc-8 export CXX/usr/bin/g-8或者暂时切换系统的 GCC 默认版本sudo update-alternatives --config gcc注意编译完成后运行时的环境不要变否则可能出现编译时用 GCC 8、运行时被其他环境变量干扰的问题。4.7 CUDA out of memory 与显存不足这个报错虽然和版本配置关系不大但用 PyTorch 1.8 的老项目尤其容易碰到RuntimeError: CUDA out of memory. Tried to allocate 512.00 MiB (GPU 0; 24.00 GiB total capacity; 22.50 GiB already allocated; ...)排查思路看有没有其他进程占用了 GPUnvidia-smi查一下如果有残留的 Python 进程kill -9掉。在你的代码开头加上import torch torch.cuda.empty_cache()检查代码是不是在循环里不断创建新的计算图导致显存碎片化。PyTorch 1.8 的显存管理机制还不像 2.x 那么智能有时候即使torch.no_grad()没写好也会爆显存。如果你的输入数据比较大尝试batch_size减半或者开启混合精度训练需要装 apex不过 PyTorch 1.8 对 apex 的兼容性有时也会出问题建议直接用 torch.cuda.amp 自带的混合精度模块。4.8 常见问题速查表报错/症状可能原因解决方案torch.cuda.is_available() False装成 CPU 版 / 驱动太低pip 卸载后重装 cu111 版升级驱动找不到 libcudart.so.11.0环境变量未配置或指向其他版本检查 LD_LIBRARY_PATH、conda 装 cudatoolkitgzip: stdin: invalid compressed data.run 文件下载不完整校验 MD5重新下载Cannot find suitable CUDA versionpip 版本太低升级 pip 到 21编译扩展报 GCC 相关错误GCC 和 CUDA 版本不兼容指定 GCC 8 或者 update-alternatives 切换CUDA out of memory其他进程占用 / 显存碎片杀掉残存进程、empty_cache、减小 batch_size安装后提示 nvcc 找不到PATH 没配置添加 CUDA bin 目录到 PATHtorch.utils.cpp_extension 报错找不到 cuda.h系统缺少 CUDA Toolkit 的头文件安装完整 CUDA Toolkit而非仅运行时5. 多项目环境的版本隔离与切换经验5.1 用 conda 环境做隔离是王道我见过很多老手在自己的主力机器上同时维护 PyTorch 1.8、1.12、2.0 三套环境互不干扰。核心思路就一句话让 conda 帮你管理 CUDA 运行时别让系统级 CUDA 来捣乱。具体做法每次新建项目环境都用 conda 显式指定 cudatoolkit 版本不要用 pip 往 base 环境里装任何深度学习的包在项目的 requirements.txt 里备注依赖的 CUDA 版本方便后来接手的人快速复盘5.2 不同 CUDA 版本的切换技巧如果你的工作流里确实需要多个 CUDA 版本比如同时维护两个老项目一个用 11.1一个用 10.2推荐用 symlink 方式管理系统级 CUDA# 查看当前版本 ll /usr/local/cuda # 切换版本 sudo rm -rf /usr/local/cuda sudo ln -s /usr/local/cuda-11.1 /usr/local/cuda虽然这不影响 conda 环境内的 cudatoolkit但当你需要编译系统级工具时这个软链接指向的版本就很重要了。记得每切换一次就要同步更新 PATH 和 LD_LIBRARY_PATH。Windows 下没有软链接这么方便主流做法是先装多个版本的 CUDA Toolkit装在不同目录如 CUDA\v10.2、CUDA\v11.1在“环境变量”里把 CUDA_PATH 指到你当前要用的版本用的时候手动改环境变量或者用脚本批量切换Windows 切换环境变量非常麻烦所以我强烈建议 Windows 用户尽量把 CUDA 运行时交给 conda 管理系统级只装一个最新版本即可。5.3 项目迁移时的注意事项如果是从其他机器或服务器迁移一个 PyTorch 1.8 项目过来别急着装环境。先检查这几个东西原项目的依赖清单pip freeze requirements.txt或者conda env export environment.yml显卡型号是否变了不同型号的 Compute Capability 会影响编译选项CUDA 版本是否一致如果原项目编译过自定义扩展迁移到另一台机器上必须重新编译尤其是第 3 点很多人项目中包含 .so 或 .pyd 这类编译好的二进制文件这些文件绑定了原机器的 CUDA 版本和 GPU 架构直接复制到新机器上几乎必然报错类似RuntimeError: CUDA error: no kernel image is available for execution on the device或者AssertionError: Torch not compiled with CUDA enabled遇到这类报错删除所有编译缓存如 build 目录、pycache目录重新编译扩展。6. 我踩过最大的坑与最后的建议6.1 在 Windows 上装 CUDA 11.1 时的 Visual Studio 联动问题如果你在 Windows 上做 CUDA 开发一定会碰到 Visual Studio 集成的问题。CUDA 11.1 需要 Visual Studio 2019 的 16.x 版本支持但很多人的 VS 已经升级到了 2022。安装 CUDA 时如果检测不到匹配的 VS 版本会提示CUDA Visual Studio Integration Unsupported: No supported version of Visual Studio was found这个报错仅影响你用 VS IDE 做 CUDA C 调试不影响 PyTorch 的运行。如果你不写自定义 CUDA 算子完全可以忽略这个警告。如果要解决办法有两个安装 VS 2019 并保留或者安装 VS 2022 的旧版 Build ToolsCUDA 11.1 不支持 VS 2022必须用 2019。让 CUDA 在安装时跳过 VS 集成检查安装 exe 时选择“自定义”取消勾选 “Visual Studio Integration” 组件。6.2 不要盲目追求最新版本这几年我发现一个规律在深度学习环境配置这件事上“最新”不等于“最好”反而“刚好能用”才是最大幸福。PyTorch 1.8 虽然不是新版本但它的稳定性和社区积累足够厚遇到任何问题都能在 GitHub Issues 或 StackOverflow 上找到现成答案。反观一些刚发布的新版本装完就遇到新 bug连报错都搜不到几个帖子那种无助感才是真的崩溃。当然如果你的项目是全新启动没有历史包袱我建议直接用 PyTorch 2.x CUDA 11.8 或 12.x性能更好代码也不用改太多。但如果是复现一个老论文、跑一个教程里的代码或者维护一个生产系统那 CUDA 11.1 PyTorch 1.8 这套组合依然能打按照本文的步骤配置一遍踩坑的概率会小很多。6.3 最后分享一个实用小技巧其实配置完环境之后我非常建议你做一个环境自检脚本把 torch、CUDA、cuDNN、GPU 信息、Python 版本、常用扩展是否可编译全部打印出来存成env_check.pyimport platform import subprocess import sys import torch print( * 60) print(Python version:, sys.version.split()[0]) print(Operating system:, platform.platform()) print(PyTorch version:, torch.__version__) print(CUDA available:, torch.cuda.is_available()) if torch.cuda.is_available(): print(CUDA version:, torch.version.cuda) print(cuDNN version:, torch.backends.cudnn.version()) print(GPU count:, torch.cuda.device_count()) for i in range(torch.cuda.device_count()): print(fGPU {i}: {torch.cuda.get_device_name(i)}) print(f Capability: {torch.cuda.get_device_capability(i)}) print(f Memory: {torch.cuda.get_device_properties(i).total_memory / 1024**3:.1f} GB) # 简单跑一个矩阵乘法验证计算是否正确 if torch.cuda.is_available(): a torch.randn(1000, 1000, devicecuda) b torch.randn(1000, 1000, devicecuda) c torch.matmul(a, b) # 等待 CUDA 异步执行完成 torch.cuda.synchronize() print(Matrix multiplication test: PASSED) print( * 60) print(All checks completed.)把这段脚本放在项目的根目录以后每次换机器或者升级环境跑一遍就知道一切是否正常。搞深度学习的人时间都很宝贵与其每次装完环境后焦虑地盯着 torch.cuda.is_available() 的输出不如把这个检查脚本固化下来一劳永逸。说实话环境配置这件事看起来繁琐但底层的逻辑其实很清晰版本匹配、环境隔离、路径正确、显存够用。把这四件事都搞明白了无论以后 PyTorch 出到 3.x 还是 4.x你都能轻松应对。希望这篇文章能帮你少走些弯路把更多时间花在模型和代码上而不是跟环境死磕。
返回列表