💡

Vitis-AI でdockerのベースイメージを変更する(3.0)

2025/02/09に公開

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分岐が実行されます。

docker_build.sh
# 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の初めのほうに定義があるので、対応するものを探して記述します。

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