AMD深度学习入门指南：从环境搭建到首个神经网络模型

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

为什么选择 AMD 做深度学习？

近年来 AMD GPU 凭借高性价比和开源生态（ROCm/Radeon Open Compute）在深度学习领域崭露头角。与同价位 N 卡相比，像 Radeon RX 6800 这样的显卡能提供：

• 显存优势：16GB GDDR6 显存轻松应对大 batch size 训练
• 开源红利：ROCm 工具链完全开源，支持 PyTorch/TensorFlow 主流框架
• 性价比突出：同等算力下价格通常比 N 卡低 20%-30%

关键检查点：
– 确认显卡在 ROCm 支持列表（如 Radeon VII/RX 5000/6000 系列）
– 需 Linux 系统（推荐 Ubuntu 20.04 LTS）

环境搭建步步通

1. 系统准备

# 查看 GPU 信息（需 root 权限）sudo lspci | grep -i amd

2. 安装 ROCm 5.x

# 添加 ROCm 仓库（普通用户需 sudo）sudo apt update && sudo apt install -y wget
wget -q -O - https://repo.radeon.com/rocm/rocm.gpg.key | sudo apt-key add -
echo 'deb [arch=amd64] https://repo.radeon.com/rocm/apt/debian/ ubuntu main' | sudo tee /etc/apt/sources.list.d/rocm.list

# 安装核心组件
sudo apt update && sudo apt install -y rocm-hip-sdk rocm-opencl-runtime

# 验证安装
/opt/rocm/bin/rocminfo | grep 'Agent.*GPU'

3. PyTorch 环境配置

# 创建 conda 环境（需已安装 Anaconda）conda create -n amd_torch python=3.8 -y
conda activate amd_torch

# 安装 ROCm 版 PyTorch
pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/rocm5.2

关键检查点：
– 运行 python -c "import torch; print(torch.cuda.is_available())" 应返回 True
– 若报错 HSA_DEVICE_NOT_FOUND，尝试sudo usermod -a -G video $USER 后重启

第一个神经网络实战

以 MNIST 手写数字识别为例，演示 CNN 模型的 ROCm 加速效果：

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

# 检测设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")

# 数据加载（自动下载）transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
train_set = datasets.MNIST('./data', download=True, train=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_set, batch_size=64, shuffle=True)

# 定义 CNN 模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = torch.max_pool2d(torch.relu(self.conv2(x)), 2)
        x = torch.flatten(x, 1)
        x = torch.relu(self.fc1(x))
        return self.fc2(x)

model = Net().to(device)
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()

# 训练循环
for epoch in range(5):
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
    print(f'Epoch {epoch+1}, Loss: {loss.item():.4f}')

关键检查点：
– 使用 watch -n 1 rocm-smi 可监控 GPU 利用率
– batch size 建议从 64 开始逐步增加至显存占满

性能对比实测

测试环境：
– AMD RX 6800 (16GB) vs NVIDIA RTX 3070 (8GB)
– batch size=128, ResNet50 模型

常见问题解决方案

1. 内核版本不匹配

   # 查看当前内核
   uname -r
   # 安装对应头文件
   sudo apt install linux-headers-$(uname -r)

2. 权限不足报错

   # 将用户加入 video 组
   sudo usermod -a -G video $USER
   # 检查组权限
   groups | grep video

3. PyTorch 找不到 GPU

   # 确认 ROCm 版本匹配
   pip uninstall torch torchvision
   pip install torch --extra-index-url https://download.pytorch.org/whl/rocm$(/opt/rocm/bin/rocminfo | grep 'ROCm 版本' | awk '{print $2}')

进阶：CUDA 代码迁移

通过 HIP 工具可将 CUDA 代码转换为 AMD 平台代码：

1. 安装转换工具

   sudo apt install hipify-clang

2. 转换示例

   hipify-perl my_cuda_code.cu > my_hip_code.cpp

3. 编译命令差异

   # 原 CUDA 编译
   nvcc -o cuda_prog cuda_code.cu
   # HIP 编译
   hipcc -o hip_prog hip_code.cpp

写在最后

经过完整的环境搭建和模型训练测试，AMD 平台在深度学习工作流中展现出不错的竞争力。虽然生态成熟度仍在追赶阶段，但对于预算有限又需要大显存的场景（如 NLP 大模型微调），AMD GPU 确实是个高性价比选择。后续可以尝试：

• 使用 ROCProfiler 进行性能分析
• 尝试混合精度训练
• 探索 ONNX Runtime 的 ROCm 后端支持