44 lines
1.3 KiB
Docker
44 lines
1.3 KiB
Docker
# 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"]
|