Мануал: Запуск PyTorch с Radeon RX 6600 в Docker на Ubuntu 22.04

Требования

• ОС: Ubuntu Server 22.04 LTS
• GPU: AMD Radeon RX 6600 (gfx1032)
• ОЗУ: 20 ГБ (или больше)
• Диск: 50+ ГБ свободного места
• Пакеты: curl, docker.io

Шаг 1: Подготовка хоста

1.1 Установка драйверов ROCm 5.3

1. Удалите предыдущие версии ROCm:

   sudo amdgpu-install --uninstall
   sudo apt purge -y rocm-* hip-* hsa-* roc* comgr*
   sudo apt autoremove -y
   sudo apt autoclean

2. Установите ROCm 5.3:

   wget https://repo.radeon.com/amdgpu-install/5.3/ubuntu/jammy/amdgpu-install_5.3.50300-1_all.deb
   sudo apt install ./amdgpu-install_5.3.50300-1_all.deb
   sudo amdgpu-install --usecase=rocm --no-dkms
   sudo reboot

3. Проверьте установку:

   dpkg -l | grep rocm
   rocm-smi
   rocminfo | grep gfx1032

1.2 Установка Docker

4. Установите Docker:

   sudo apt update
   sudo apt install -y docker.io
   sudo systemctl enable docker
   sudo systemctl start docker

5. Проверьте:

   docker --version

Шаг 2: Создание Dockerfile

Создайте файл Dockerfile:

FROM rocm/dev-ubuntu-22.04:5.3

# Установка зависимостей
RUN apt-get update && apt-get install -y \
 python3 \
 python3-pip \
 rocm-libs \
 && rm -rf /var/lib/apt/lists/*

# Установка NumPy
RUN pip3 install --no-cache-dir --force-reinstall "numpy<2.0"

# Установка PyTorch с ROCm 5.3
RUN pip3 install --no-cache-dir --force-reinstall torch==2.0.1 torchvision==0.15.2 \
 --index-url https://download.pytorch.org/whl/rocm5.3

# Установка Flask и Werkzeug
RUN pip3 install --no-cache-dir flask werkzeug && \
 pip3 install --no-cache-dir --force-reinstall "numpy<2.0"

# Установка переменной окружения для поддержки gfx1032
ENV HSA_OVERRIDE_GFX_VERSION=10.3.0

# Копируем код
WORKDIR /app
COPY app.py .

# Запуск
CMD ["python3", "app.py"]

Шаг 3: Создание приложения Flask (app.py)

Создайте файл app.py:

from flask import Flask, request, jsonify
import torch
import torch.nn as nn

app = Flask(__name__)

model = nn.Sequential(
 nn.Linear(10, 5),
 nn.ReLU(),
 nn.Linear(5, 1)
).to('cuda')

@app.route('/predict', methods=['POST'])
def predict():
 data = request.json['input']
 input_tensor = torch.tensor(data, dtype=torch.float32).to('cuda')
 with torch.no_grad():
 output = model(input_tensor)
 return jsonify({'output': output.cpu().tolist()})

if __name__ == '__main__':
 app.run(host='0.0.0.0', port=5000)

Шаг 4: Сборка и запуск контейнера

1. Соберите образ:

   docker build -t rocm-pytorch-rx6600 .

2. Запустите контейнер:

   docker run --device=/dev/kfd --device=/dev/dri -p 5000:5000 rocm-pytorch-rx6600

Шаг 5: Тестирование

Проверка GPU

docker run --device=/dev/kfd --device=/dev/dri -it rocm-pytorch-rx6600 bash
python3 -c "import torch; print(torch.cuda.is_available()); print(torch.cuda.get_device_name(0))"

Ожидаемый вывод:

True
AMD Radeon RX 6600 (или AMD Radeon Graphics)

Тест API

curl -X POST -H "Content-Type: application/json" -d '{"input": [1,2,3,4,5,6,7,8,9,10]}' http://192.168.0.100:5000/predict

Ожидаемый вывод:

{"output": [-3.884134292602539]} # Число может варьироваться, так как модель не обучена

Примечания

• Переменная HSA_OVERRIDE_GFX_VERSION: Используется для обхода ограничений поддержки gfx1032.
• Производительность: Замените модель на свою для реальных задач.
• Обновления: Следите за новыми версиями ROCm для лучшей поддержки.