41 lines
1.1 KiB
Docker
41 lines
1.1 KiB
Docker
# YOLOv8 training image (NVIDIA GPU)
|
|
#
|
|
# Image contains Python / PyTorch / Ultralytics + train script.
|
|
# Dataset (*.pt weights, /data/facepp, runs/) are mounted at runtime.
|
|
|
|
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-pip \
|
|
&& rm -rf /var/lib/apt/lists/*
|
|
|
|
# PyTorch 2.5.1 + torchvision 0.20.1 (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
|
|
|
|
COPY train_yolov8.py data.yaml ./
|
|
|
|
ENTRYPOINT ["python3", "train_yolov8.py"]
|
|
CMD ["--model", "yolov8n.pt", "--device", "0"]
|