如何设置PyTorch在AMD GPU上运行 CUDA设备选择命令的等效写法是什么

PyTorch AMD ROCm 平台设备指定等效写法 #

适配ROCm的PyTorch版本已经对CUDA相关API做了兼容映射，你原来的CUDA判断代码几乎不需要做额外修改：

• 优先使用官方兼容写法（原有CUDA代码无需修改即可直接跑在AMD GPU上）：

# 安装ROCm版本PyTorch后，torch.cuda相关API会自动映射到AMD GPU底层驱动
torch.device('cuda' if torch.cuda.is_available() else 'cpu')

• 如果你需要显式指定ROCm原生设备标识，可以用如下写法：

# torch.has_hip会直接检测当前环境是否支持ROCm AMD GPU
torch.device('hip' if torch.has_hip else 'cpu')

注意：以上写法生效的前提是你已经安装了和当前环境ROCm版本匹配的PyTorch发行版，若安装的是CPU版或CUDA版PyTorch，无法识别AMD GPU。如果需要强制指定某一块AMD GPU，和CUDA逻辑完全一致，写cuda:0/hip:0即可指定序号为0的显卡。

内容的提问来源于stack exchange，提问作者southernKid33