Skip to content

AMD平台在Ubuntu安装ROCm并使用PyTorch

ROCm（Radeon Open Compute）是继 CUDA 之后的第二大 GPU 并行计算平台，能够为 PyTorch 提供加速支持。以下是如何在 AMD GPU 上使用 ROCm 加速 PyTorch 的简要步骤：

硬件/系统配置

  • 硬件：以 AMD Radeon RX 7800 XT + AMD Ryzen R5 9600X + 32GB DDR5 的配置为例，具体需求请参考 AMD 官方文档。
  • 系统：AMD 官方支持 Ubuntu、Red Hat Enterprise Linux (RHEL)、SUSE Linux Enterprise Server (SLES)，其他相关发行版如 Linux Mint、Rocky Linux、OpenSUSE 等大概率也可正常使用。ROCm 暂不支持 Windows，需借助 WSL2。

ROCm 安装

在 Ubuntu 上，执行以下命令完成安装：

sudo apt update
sudo apt install "linux-headers-$(uname -r)" "linux-modules-extra-$(uname -r)"
sudo usermod -a -G render,video $LOGNAME
wget https://repo.radeon.com/amdgpu-install/6.2.2/ubuntu/jammy/amdgpu-install_6.2.60202-1_all.deb
sudo apt install ./amdgpu-install_6.2.60202-1_all.deb
sudo apt update
sudo apt install amdgpu-dkms rocm

配置 ld

sudo tee --append /etc/ld.so.conf.d/rocm.conf <<EOF
/opt/rocm/lib
/opt/rocm/lib64
EOF
sudo ldconfig

将 ROCm 的可执行文件添加至 PATH

可选择以下任一方法：

  • 使用 update-alternatives：
update-alternatives --list rocm
update-alternatives --config rocm
  • 使用 environment-modules：
module avail
module load rocm/<version>
  • 手动配置：在 .bashrc 中添加 export PATH=$PATH:/opt/rocm-6.2.2/bin

验证安装状态

dkms status
rocminfo
clinfo
apt list --installed

重启

reboot

PyTorch 安装

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm6.2

验证 PyTorch 及 ROCm

在 Python 中执行以下指令：

import torch
x = torch.rand(5, 3)
print(x)
print(torch.cuda.is_available())

若 torch.cuda.is_available() 返回 True，则表示 ROCm 已成功启用。如遇问题，可查看日志并尝试重新添加用户至 render 和 video 组后重启系统。

PyTorch ROCm 现已发布 – 如何选择 Radeon GPU 作为设备

什么是 PyTorch ROCm？

PyTorch ROCm 是 PyTorch 的一个扩展，基于 ROCm（Radeon 开放计算平台）提供对 AMD Radeon GPU 的支持。ROCm 是一个开源的软件平台，用于使用 AMD GPU 进行高性能计算和机器学习。与使用 CUDA 的 PyTorch 相比，PyTorch ROCm 可以在某些情况下提供类似的性能，甚至在特定任务上表现更优。

如何选择 Radeon GPU 作为设备

步骤 1：安装 ROCm

首先，需要安装 ROCm。请访问 AMD 官方网站以获取最新版本的 ROCm，并根据其提供的说明进行安装。

步骤 2：安装 PyTorch 和 PyTorch ROCm 扩展

安装了 ROCm 之后，需要安装 PyTorch 和 PyTorch ROCm 扩展。可以通过以下命令使用 conda 进行安装：

conda install pytorch torchvision torchaudio rocm -c pytorch -c rapidsai
步骤 3：选择 Radeon GPU 作为设备

要在 PyTorch 中选择 Radeon GPU 作为设备，需要通过以下代码进行设置：

import torch
device = torch.device("hip:0")  # 选择第一个 Radeon GPU
# 如果有多个 Radeon GPU，可以选择具体的 GPU 设备，例如：
# device = torch.device("hip:1")  # 选择第二个 Radeon GPU
torch.set_default_tensor_type(torch.HalfTensor)

示例：在 PyTorch 中利用 Radeon GPU 进行深度学习训练

以下是一个示例，展示了如何在 PyTorch 中利用 Radeon GPU 进行深度学习训练：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

# 选择 Radeon GPU 作为设备
device = torch.device("hip:0")

# 定义一个简单的卷积神经网络
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.fc1 = nn.Linear(32 * 26 * 26, 10)

    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = x.view(-1, 32 * 26 * 26)
        x = self.fc1(x)
        return x

# 创建示例数据
x_train = torch.randn(64, 1, 28, 28).to(device)
y_train = torch.randint(0, 10, (64,)).to(device)
dataset = TensorDataset(x_train, y_train)
dataloader = DataLoader(dataset, batch_size=32)

# 初始化模型、损失函数和优化器
model = SimpleCNN().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 训练模型
for epoch in range(10):
    for batch_x, batch_y in dataloader:
        optimizer.zero_grad()
        output = model(batch_x)
        loss = criterion(output, batch_y)
        loss.backward()
        optimizer.step()

# 测试模型
with torch.no_grad():
    test_x = torch.randn(10, 1, 28, 28).to(device)
    output = model(test_x)
    print(output)

总结

本文介绍了如何在 PyTorch 中选择 Radeon GPU 作为设备，并展示了在 PyTorch 中利用 Radeon GPU 进行深度学习训练的示例。PyTorch ROCm 的发布为 AMD 用户提供了更多选择，让他们能够充分发挥 Radeon GPU 的潜力，并在机器学习和高性能计算领域取得更好的性能。希望本文对您有所帮助！

Ref

https://geek-docs.com/pytorch/pytorch-questions/74_pytorch_pytorch_rocm_is_out_how_to_select_radeon_gpu_as_device.html#google_vignette https://dkrain.com/posts/rocm%E7%AE%80%E5%8D%95%E5%85%A5%E9%97%A8-%E4%BD%BF%E7%94%A8amd%E6%98%BE%E5%8D%A1%E5%8A%A0%E9%80%9Fpytorch/#google_vignette