moophlo/comfyui Docker 镜像 - 轩辕镜像 | Docker 镜像高效稳定拉取服务

AMD GPU PyTorch镜像文档

镜像概述

本Docker镜像专为AMD GPU环境设计，基于pytorch/rocm6.2.4构建，提供了在AMD GPU上运行PyTorch应用的基础环境。它整合了ROCm 6.2.4计算栈与PyTorch框架，确保在AMD GPU硬件上实现高效的深度学习计算。

核心功能与特性

• AMD GPU优化：针对AMD GPU架构优化，充分利用ROCm生态系统
• 版本锁定：基于PyTorch ROCm 6.2.4版本，确保环境稳定性和兼容性
• 即开即用：预配置PyTorch与ROCm依赖，无需手动安装底层驱动和库

使用场景

• AMD GPU环境下的深度学习模型训练
• 基于PyTorch的机器学习推理任务
• 依赖ROCm加速的科学计算应用开发与测试

使用方法

基本运行命令

bash
docker run --rm -it \
  --device=/dev/kfd \
  --device=/dev/dri \
  --group-add video \
  <镜像名称>

关键参数说明

• --device=/dev/kfd：挂载AMD GPU设备节点，用于ROCm运行时访问
• --device=/dev/dri：挂载显示设备接口，支持GPU显示输出
• --group-add video：添加video用户组权限，确保容器内对GPU设备的访问权限

典型应用示例

启动交互式PyTorch环境：

bash
docker run --rm -it \
  --device=/dev/kfd \
  --device=/dev/dri \
  --group-add video \
  -v $(pwd):/workspace \
  <镜像名称> python

在容器内验证GPU可用性：

python
import torch
print(torch.cuda.is_available())  # 应输出True
print(torch.cuda.get_device_name(0))  # 显示AMD GPU型号

注意事项

• 宿主机需已安装ROCm驱动，且版本与镜像兼容（建议ROCm 6.2.x系列）
• 运行时需确保容器拥有足够的GPU访问权限
• 对于多GPU环境，可通过--device参数指定多个GPU设备节点