2条命令跑通AMD GPU：TensorFlow的ROCm实测

在AMD GPU上跑TensorFlow，过去最劝退的一步是环境配置。ROCm是AMD的GPU计算平台，目标是在AMD显卡上提供高性能计算和机器学习能力。而一个支持ROCm的TensorFlow容器，本质上是把依赖关系预先打包好的可移植环境——TensorFlow已经针对AMD GPU做过优化，你不需要手动装一堆东西、调一堆版本。

这篇指南要做的事很具体：在一台带AMD GPU的服务器上，下载并运行一个支持ROCm的TensorFlow容器，然后用pip把TensorFlow直接装到宿主机上，走ROCm计算平台。做完之后，你会得到两个可用的TensorFlow GPU加速环境——一个在容器里，一个原生装在宿主机上。

先确认前置条件

开始之前有两件事必须先就位。第一，一台Ubuntu 24.04服务器，并且配备支持ROCm的AMD GPU。第二，Docker已经安装并配置好ROCm GPU支持。这两项缺一不可，尤其是Docker那部分——如果Docker本身没有配好ROCm GPU支持，后面的容器跑起来也访问不到显卡。

容器路线：拉镜像、跑起来、验显卡

第一步是拉取支持ROCm的TensorFlow容器镜像：

docker pull rocm/tensorflow:latest

镜像到手之后，用一条docker run命令启动一个临时容器。这条命令里几个参数值得留意：--device=/dev/kfd 和 --device=/dev/dri 负责把GPU设备暴露给容器，这是ROCm工作负载能用到显卡的前提；--security-opt seccomp=unconfined 放开安全配置限制；--shm-size 8G 把共享内存调到8G。加上--rm和-it，容器用完即走。

docker run --rm -it --device=/dev/kfd --device=/dev/dri --security-opt seccomp=unconfined --shm-size 8G rocm/tensorflow:latest

进到容器里之后，先跑 rocm-smi 看显卡状态，再用一行Python确认TensorFlow能不能识别到设备：

python3 -c 'import tensorflow; print(tensorflow.config.list_physical_devices())'

输出应该会把所有设备连同规格列出来。看到设备列表，说明容器这条路走通了。确认完直接 exit 退出，临时容器会被销毁，不留痕迹。

宿主机路线：版本匹配是关键

容器适合快速验证，但如果想让TensorFlow直接跑在宿主机上，就得走pip安装这条路。这里最容易踩坑的地方是版本对应关系。

先查ROCm版本：

amd-smi version

拿到版本号之后，去官方安装页面查支持矩阵，找到与当前ROCm版本匹配的TensorFlow版本。指南给了一个具体例子：如果ROCm版本是6.2.x，那么兼容的TensorFlow版本是2.14、2.15、2.16。

然后用预编译的wheel包安装：

python3 -m pip install tensorflow-rocm=={TENSORFLOW_VERSION} -f https://repo.radeon.com/rocm/manylinux/rocm-rel-{ROCM_VERSION}

这条命令里有两个占位符要替换。{TENSORFLOW_VERSION} 填上一步从支持矩阵查到的版本号；{ROCM_VERSION} 只取ROCm版本号的前两位数字。

还是拿ROCm 6.2.x举例，替换之后的完整命令长这样：

python3 -m pip install tensorflow-rocm==2.16 -f https://repo.radeon.com/rocm/manylinux/rocm-rel-6.2

装完之后，同样用一行Python验证GPU是否可用。如果设备列表里出现了AMD显卡，说明宿主机这条路也走通了。至此，容器和宿主机两个TensorFlow GPU加速环境都已就绪。