pytorch + AMD卡 (docker快速构建A卡pytorch环境, 5分钟不折腾版)

74 阅读 0 评论 49 点赞

我是靠谱客的博主寂寞毛豆，最近开发中收集的这篇文章主要介绍pytorch + AMD卡 (docker快速构建A卡pytorch环境, 5分钟不折腾版)，觉得挺不错的，现在分享给大家，希望可以做个参考。

概述

环境

宿主机环境
ubuntu20.02(其实能跑docker的linux都可以)
CPU: AMD Ryzen 9 5950X
GPU: AMD 6700XT 12G

理论上可以支持绝大部分的A卡包含570 580
容器环境
ROCm5.0
pytorch1.10.0
python3.7

启动容器

docker run 
    --device=/dev/kfd 
    --device=/dev/dri 
    --cap-add=SYS_PTRACE 
    --security-opt seccomp=unconfined 
    --ipc=host 
    -dit 
    --rm 
    -e  HSA_OVERRIDE_GFX_VERSION=10.3.0 
    amdih/pytorch:rocm5.0_ubuntu18.04_py3.7_pytorch_1.10.0

在这里插入图片描述

验证是否可用

docker exec -it 容器id bash
# 进入容器后输入一下命令
python
import torch
print(torch.cuda.is_available())

在这里插入图片描述

快速部署ssh直通版(ssh =直达=> 容器)

# 直接通过 exit run 的命令
# 并添加:2022端口ssh访问容器
# 设置root的密码为root
docker exec -it 
  $(docker run 
    --device=/dev/kfd 
    --device=/dev/dri 
    --cap-add=SYS_PTRACE 
    --security-opt seccomp=unconfined 
    --ipc=host 
    -dit 
    --rm 
    -p 2022:22 
    -e  HSA_OVERRIDE_GFX_VERSION=10.3.0 
    amdih/pytorch:rocm5.0_ubuntu18.04_py3.7_pytorch_1.10.0) 
  bash -c 
	"apt-get update && 
	apt-get install openssh-server -y && 
	echo 'PermitRootLogin yes' >> /etc/ssh/sshd_config && 
	echo 'PasswordAuthentication yes' >> /etc/ssh/sshd_config && 
	/etc/init.d/ssh start && 
	echo 'root:root' | chpasswd"

在这里插入图片描述

验证1, docker exec验证

docker exec -it 容器ID python
import torch
print(torch.cuda.is_available())

在这里插入图片描述

验证2, ssh 验证

# 这里的2022是上面的 -p映射的 
ssh root@容器宿主机IP -p2022

# 这里需要输入环境变量运行python
HSA_OVERRIDE_GFX_VERSION=10.3.0 python
import torch
print(torch.cuda.is_available())

在这里插入图片描述