Vitis-AI でdockerのベースイメージを変更する(3.0)
Viti-AI 3.0 使用の際、デフォルトのgpu向けdocker_build.shでは、nvccが使えないなど何かと不便です。また、gpuの世代によっては合わないこともあるので、変更していきます。
一例として、pytorchで環境構築をする例をとります。
※ vitis-aiの2025/1/1時点での最新バージョンは 3.5なので、わざわざ3.0を使う方向けの備忘録です。
build.shの書き換え
もともとのコンテナ構築手順は下記です。
nvidiaのgpuを利用する上では、下記の関数のif分岐が実行されます。
# Execute
function execute
{
add_args=""
echo "SKIP:${SKIP_BUILD_BASE_IMAGE}, docker type:${DOCKER_TYPE}\n"
if [[ "$SKIP_BUILD_BASE_IMAGE" == "0" ]];then
if [[ "$DOCKER_TYPE" == 'cpu' ]];then
VAI_BASE="ubuntu:20.04"
fi
if [[ "$DOCKER_TYPE" == 'gpu' ]];then
VAI_BASE="nvidia/cuda:11.3.1-cudnn8-runtime-ubuntu20.04"
# 以下略
VAI_BASE
がベースとなるイメージですね。名前に「runtime」と入るものだと、nvccなどが使えない場合があるようです。逆に「devel」だと使えるとか。
nvidiaが現段階で提供しているコンテナイメージをNGC Catalogから探して、IDを入れてあげると、それがダウンロードされるつくりの様です。
仮にcudnn8, cuda11.1系のubuntu20.04コンテナを使うとすると、
VAI_BASE="nvcr.io/nvidia/cuda:11.1.1-cudnn8-devel-ubuntu20.04"
のようになりますね。
ubuntuのバージョンが20.04以外の場合、XRT_URL, XRM_URLも変えてあげるのが無難でしょう。
docker_build.shの初めのほうに定義があるので、対応するものを探して記述します。
XRT_URL="${XRT_URL:-https://www.xilinx.com/bin/public/openDownload?filename=XXXXXX-amd64-xrt.deb}"
XRM_URL="${XRM_URL:-https://www.xilinx.com/bin/public/openDownload?filename=XXXXXX-x86_64.deb}"
```
docker imageの名称を独自に設定する場合は、`IMAGE_LATEST_TAG`変数などの値をいじるとできそうです。
## pytorchのバージョンを変更する
上記によってコンテナを起動できるようになったと仮定して。
pytorchのバージョンをデフォルトのもの以外に変更する場合、リポジトリ内にある `docker/common/replace_pytorch.sh`が利用できるとのことです。shの中身を書き換えれば1.9, 1.7などにできそうですね。
このシェルがうまく実行できない場合、手動で入れ替えることも一応可能です。起動したシェルで以下を実行します。
```sh
conda rm torch
pip uninstall torch
pip uninstall torchvision
# その後インストールしたいpytorchをインストールする
pytorchのインストール方法は公式ページをみながら行いました。
はまったところ
pipなどでpytorchをインストールする場合、ホストマシンに搭載のGPUのアーキテクチャタグ(sm86とか)によっては非対応となっている場合がありました。コンテナで少々古いpytorchを利用した際に踏んでしましました。
場合によっては最初から入れ直しになるかも。
Discussion