Zenn
📑

NVIDIA Docker を用いた Pytorchの環境構築の備忘録

2024/03/27に公開
Docker
PyTorch
tech

以前に，Geforce のドライバ，NVIDIA Docker と cuda はインストール済みなので，割愛．

動作の確認

すでにインストール済みの NVIDIA Docker が正常に動作するかの確認をするべきなので，以下のコマンドで確認．

# ドライバの動作の確認
$ nvidia-smi
# NVIDIA Docker の動作の確認
$ docker run --gpus all --rm nvidia/cuda:12.2.0-base-ubuntu22.04 nvidia-smi

Dockerfile の作成

pytorch/pytorch:2.2.1-cuda12.1-cudnn8-devel を使うのが最も簡単だった．

FROM pytorch/pytorch:2.2.1-cuda12.1-cudnn8-devel

GPUが認識されているか確認

まず，コマンドでイメージとコンテナを作成する．

$ docker build -t pytroch .
$ docker run -it --gpus all --rm pytorch bash

つぎに，pytrochからGPUが認識されているか確認する．

root@898aea227346:/workspace# python3
Python 3.10.13 (main, Sep 11 2023, 13:44:35) [GCC 11.2.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> torch.cuda.is_available()
True
>>> torch.cuda.device_count()
1

devcontainerの設定

devcontainer.json の設定は以下のようにした．
--gpus all オプションを設定してます．

{
    "name": "NVIDIA Docker Ubuntu 22.04 dev",
    "build": {
        "dockerfile": "Dockerfile"
    },
    "runArgs": [
        "--gpus",
        "all"
    ],
    "remoteUser": "vscode",
    "customizations": {
        // Configure properties specific to VS Code.
        "vscode": {
            // Set *default* container specific settings.json values on container create.
            "settings": {
                "terminal.integrated.defaultProfile.linux": "bash",
                "editor.formatOnSave": true
            },
            "extensions": [
                "ms-python.python"
            ]
        }
    },
    "features": {
        "common": {
            "username": "vscode",
            "uid": "automatic",
            "gid": "automatic",
            "installZsh": true,
            "installOhMyZsh": false,
            "upgradePackages": false,
            "nonFreePackages": false
        }
    }
}

Discussion