pytorch + GPUのコンテナ環境設定（マルチステージビルドを添えて）

やりたかったこと

Dockerのマルチステージビルドを用いて、GPU付きマシンでpytorch学習環境を作ってみました。
下記ファイルを同ディレクトリに配置してビルドすれば通ります。

.env

UID=
GID=
# bash idから取得

docker-compose.yml

version: '3.8'

services:
  app-name:
    build:
      context: .
      dockerfile: Dockerfile
      args:
      - USER=$USER
      - UID=$UID
      - GID=$GID
    env_file:
      - .env
    container_name: container
    runtime: nvidia
    shm_size: 64gb
    environment:
      - NVIDIA_VISIBLE_DEVICES=1
      - NVIDIA_DRIVER_CAPABILITIES=compute,utility
    volumes:
      - .:/workspace:cached
    tty: true
    restart: always

Dockerfile

# ベースイメージとして公式のCUDAイメージを使用
FROM nvidia/cuda:11.3.1-cudnn8-runtime-ubuntu20.04 as base
ENV DEBIAN_FRONTEND noninteractive

# 作業ディレクトリを設定
WORKDIR /workspace

# 必要なライブラリをインストール
RUN apt-get update && apt-get install -y --no-install-recommends \
    python3-pip \
    python3-dev \
    libgl1-mesa-glx \
    libnuma1 \
    && rm -rf /var/lib/apt/lists/*
    # libgl1-mesa-glx libnuma1 はCUDAのインストール時に必要、segfaultを防ぐために必要
FROM base as dev
ARG USER=user-name-goes-here
ARG UID=1000
ARG GID=$UID
ENV TZ=Asia/Tokyo

# Create the user
RUN groupadd --gid $GID $USER \
    && useradd --uid $UID --gid $GID -m $USER --shell /bin/bash \
    # Optional: Add sudo support
    && apt-get update \
    && apt-get install -y --no-install-recommends sudo procps gdb \
    && echo $USER ALL=\(root\) NOPASSWD:ALL > /etc/sudoers.d/$USER \
    && chmod 0440 /etc/sudoers.d/$USER \
    && chown $UID:$GID -R /workspace

# Pythonパッケージをインストール
COPY requirements.txt /workspace/
RUN pip3 install --upgrade pip && \
    pip3 install -r /workspace/requirements.txt

# リポジトリの内容をコンテナ内にコピー
COPY . /workspace/

# Set the default user
USER $USER

requirements.txt

numpy
opencv-python
matplotlib
seaborn
pandas
torch
torchvision
scikit-learn