AMD显卡跑深度学习的实战指南：从ROCm环境搭建到性能调优

1次阅读

共计 2247 个字符，预计需要花费 6 分钟才能阅读完成。

近年来，NVIDIA 凭借 CUDA 生态的垄断地位，使得其显卡价格居高不下，尤其是高性能计算卡如 A100/H100，对于中小团队和个人开发者来说，采购成本成为了进入深度学习领域的门槛。而 AMD 显卡在性价比上具有明显优势，但长期以来由于缺乏成熟的深度学习支持，导致开发者望而却步。ROCm 平台的推出为 AMD 显卡在深度学习领域的应用提供了可能。

特性	ROCm	CUDA
显存管理	HIP 内存模型	CUDA 内存模型
核函数调度	HIP 内核	CUDA 内核
编程语言支持	C++, Python	C++, Python, Fortran
跨平台支持	Linux	Linux, Windows
性能优化工具	ROCProfiler	NVIDIA Nsight
深度学习框架支持	PyTorch, TensorFlow (部分)	PyTorch, TensorFlow, MXNet 等

系统准备 ：确保系统为 Ubuntu 22.04，并更新到最新版本。

sudo apt update && sudo apt upgrade -y

安装 ROCm：添加 ROCm 仓库并安装核心组件。

sudo apt install -y rocm-opencl-runtime

验证安装 ：检查 ROCm 是否安装成功。

rocminfo

解决驱动冲突 ：如果遇到驱动冲突，可以尝试以下命令。

sudo apt purge nvidia* && sudo apt autoremove

HIPIFY 是一个将 CUDA 代码转换为 HIP 代码的工具，适用于 PyTorch 模型。以下是一个简单的转换示例：

# 原始 CUDA 代码
import torch

def cuda_kernel():
    # CUDA 核函数
    pass

# 转换后的 HIP 代码
import torch

def hip_kernel():
    # HIP 核函数
    pass

TensorFlow 官方提供了支持 ROCm 的 Docker 镜像，可以直接拉取使用：

docker pull rocm/tensorflow:latest

以下是一个使用 HIP 核函数的 MNIST 训练示例：

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms

# 定义模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(784, 512)
        self.fc2 = nn.Linear(512, 10)

    def forward(self, x):
        x = x.view(-1, 784)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 初始化模型和优化器
model = Net().to('hip')
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 数据加载
train_loader = torch.utils.data.DataLoader(
    datasets.MNIST('../data', train=True, download=True,
                   transform=transforms.Compose([transforms.ToTensor(),
                       transforms.Normalize((0.1307,), (0.3081,))
                   ])),
    batch_size=64, shuffle=True)

# 训练循环
for epoch in range(10):
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to('hip'), target.to('hip')
        optimizer.zero_grad()
        output = model(data)
        loss = nn.functional.cross_entropy(output, target)
        loss.backward()
        optimizer.step()

在 PyTorch 2.2 和 ROCm 5.7 环境下，RX 7900 XTX 与 RTX 4090 在 ResNet50 上的吞吐量对比如下：