1

2026-08-18 09:32:14 +08:00
parent 66ac6f8a39
commit 725777d860
3 changed files with 71 additions and 0 deletions
@@ -0,0 +1,43 @@
+# YOLOv8 training environment (NVIDIA GPU)
+#
+# The image contains Python, PyTorch and Ultralytics only. The project
+# directory (including data.yaml, train/valid/test and *.pt files) is mounted
+# at /workspace when the container is started; this keeps the image small.
+
+FROM nvidia/cuda:12.4.1-cudnn-runtime-ubuntu22.04
+
+ENV DEBIAN_FRONTEND=noninteractive \
+    PYTHONDONTWRITEBYTECODE=1 \
+    PYTHONUNBUFFERED=1 \
+    PIP_NO_CACHE_DIR=1 \
+    NVIDIA_VISIBLE_DEVICES=all \
+    NVIDIA_DRIVER_CAPABILITIES=compute,utility
+
+RUN apt-get update \
+    && apt-get install -y --no-install-recommends \
+        ca-certificates \
+        libglib2.0-0 \
+        libgl1 \
+        libsm6 \
+        libxext6 \
+        libxrender1 \
+        python3 \
+        python3-dev \
+        python3-pip \
+    && rm -rf /var/lib/apt/lists/*
+
+# PyTorch 2.5.1 is paired with torchvision 0.20.1 and CUDA 12.4.
+RUN python3 -m pip install --upgrade pip \
+    && python3 -m pip install \
+        torch==2.5.1 \
+        torchvision==0.20.1 \
+        --index-url https://download.pytorch.org/whl/cu124 \
+    && python3 -m pip install ultralytics==8.3.0
+
+WORKDIR /workspace
+
+# These files make the image usable for a quick smoke test. The normal run
+# command mounts the whole project over /workspace.
+COPY train_yolov8.py data.yaml ./
+
+ENTRYPOINT ["python3", "train_yolov8.py"]