手里有一块AMD显卡,想跑PyTorch,第一反应往往是"能不能用"。答案是能,前提是把ROCm这套东西装对。ROCm是AMD的GPU计算平台,专门用来在AMD显卡上做高性能计算和机器学习。PyTorch本身是开源框架,用来搭建和训练机器学习模型,尤其是深度神经网络。问题出在依赖配置上——手动装一遍ROCm版的PyTorch,环境变量、驱动版本、编译选项,任何一环出错都会卡住。所以更省事的做法是先跑一个预构建的容器,ROCm版的PyTorch容器把PyTorch和依赖都配好了,开箱即用,不用自己折腾。
这篇走两条路:先在容器里把GPU加速跑通,再回到宿主机上用Pip直接装一遍,两边都验证。动手之前,服务器上得先装好Docker,并且开启ROCm的GPU支持,这一步没做,后面的命令都会失败。
容器路线:拉镜像、跑起来、验GPU
第一步,拉取支持ROCm的PyTorch容器:docker pull rocm/pytorch:latest。第二步,跑一个临时容器,这条命令的关键在于把GPU设备暴露给容器:docker run --rm -it --device=/dev/kfd --device=/dev/dri --security-opt seccomp=unconfined --shm-size 8G rocm/pytorch:latest。其中 /dev/kfd 和 /dev/dri 是ROCm工作负载必须访问的两个GPU设备节点,少了它们,容器里看不到显卡。
第三步,在容器内验证GPU是否可用,跑两条命令:rocm-smi 和 python3 -c 'import torch; print(torch.cuda.is_available())'。第一条会列出所有可用设备及其规格信息,第二条如果GPU可用,会打印 true。看到true,说明容器这条路通了。第四步,退出并销毁这个临时容器:exit。因为启动时带了 --rm 参数,退出即清理,不会在机器上留下垃圾。
宿主机路线:用Pip直接装
容器适合快速验证,但如果要在宿主机上长期开发,直接装一份更顺手。先去PyTorch的Start Locally页面选安装选项,按这个组合来:PyTorch Build选Stable,Your OS选Linux,Package选Pip,Language选Python,Compute Platform选ROCm。页面会生成一条对应的安装命令,复制到终端里执行就行。
下载过程中如果遇到超时,有两个处理办法:给安装命令加上 --timeout 参数并设一个合适的值;或者加 --no-cache-dir 参数,忽略掉不完整或损坏的缓存文件,从头重新下载。装完之后同样验证一遍:python3 -c 'import torch; print(torch.cuda.is_available())'。GPU可用的话,输出同样是 true。
两条路走完,得到什么
容器路线的价值在于隔离和可移植。一个预构建的镜像把PyTorch和ROCm的依赖关系固定下来,换台机器拉同一个镜像,行为一致,适合做验证、跑一次性任务,或者不想污染宿主机环境的时候。宿主机路线的价值在于日常开发,装好之后不用每次进容器,编辑器、调试器、本地文件系统都能直接用。两条路最后都指向同一个验证动作:torch.cuda.is_available() 返回true。这个函数名里带的是cuda,但在ROCm平台上它检测的是AMD GPU——这是PyTorch为了兼容性保留的命名,别被名字骗了。
热门跟贴