pytorchlab/pytorch
Minimal PyTorch base images for GPU (CUDA 12.6) & CPU — with optional pre-bundled YOLO weights
1.7K
pytorchlab/pytorch repository overview
pytorchlab/pytorch
Minimal, production-ready PyTorch base images for GPU and CPU workloads — with optional pre-bundled YOLO weights for offline inference.
Overview
pytorchlab/pytorch provides slim, multi-variant Docker images on top of the official pytorch/pytorch base. Each image adds only what is needed for real-world deployments:
- System libraries required by OpenCV (
libGL,libglib2.0,libpq,libgeos) - Clean layer —
__pycache__removed at build time to minimize image size - YOLO variants — pre-downloaded Ultralytics weights bundled inside the image so that ECS/Kubernetes workers have zero outbound dependency at runtime
Available Tags
GPU (CUDA 12.6 + cuDNN 9)
Built from
pytorch/pytorch:2.9.1-cuda12.6-cudnn9-runtime. Requires NVIDIA drivers compatible with CUDA 12.6+.
| Tag | Description |
|---|---|
2.9.1-cuda12.6-cudnn9-runtime | PyTorch + CUDA + system libs. Generic GPU base. |
2.9.1-cuda12.6-cudnn9-yolo26m | Above + YOLO v2.6m & v2.6n weights in /app/weights (Ultralytics offline). |
2.9.1-cuda12.6-cudnn9-yolo11m | Above + YOLO v11m & v2.6n weights in /app/weights (Ultralytics offline). |
CPU (Python 3.11 Slim)
Built from
python:3.11-slimwith CPU-only PyTorch (~1.5 GB lighter than the CUDA variant).
| Tag | Description |
|---|---|
2.4.1-cpu-py3.11-slim | PyTorch (CPU-only) + system libs. Generic CPU base. |
cpu-python3.11-yolo26m | Above + YOLO v2.6m & v2.6n weights in /app/weights (Ultralytics offline). |
cpu-python3.11-yolo11m | Above + YOLO v11m & v2.6n weights in /app/weights (Ultralytics offline). |
Quick Start
# GPU — verify CUDA is available
docker pull pytorchlab/pytorch:2.9.1-cuda12.6-cudnn9-runtime
docker run --gpus all --rm pytorchlab/pytorch:2.9.1-cuda12.6-cudnn9-runtime \
python -c "import torch; print('CUDA:', torch.cuda.is_available())"
# CPU — lightweight inference
docker pull pytorchlab/pytorch:2.4.1-cpu-py3.11-slim
docker run --rm pytorchlab/pytorch:2.4.1-cpu-py3.11-slim \
python -c "import torch; print('PyTorch:', torch.__version__)"
Use as a Base Image
GPU workload
FROM pytorchlab/pytorch:2.9.1-cuda12.6-cudnn9-runtime
WORKDIR /app
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
COPY . .
CMD ["python", "train.py"]
CPU workload / lightweight inference
FROM pytorchlab/pytorch:2.4.1-cpu-py3.11-slim
WORKDIR /app
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
COPY . .
CMD ["python", "inference.py"]
Offline YOLO inference (no internet required at runtime)
FROM pytorchlab/pytorch:2.9.1-cuda12.6-cudnn9-yolo26m
# Weights are pre-downloaded at /app/weights/yolo26m.pt & yolo26n.pt
# YOLO_OFFLINE=1 and YOLO_CONFIG_DIR=/app/weights are already set
WORKDIR /app
COPY . .
CMD ["python", "detect.py"]
Requirements
| Requirement | Notes |
|---|---|
| Docker Engine 20.10+ | Tested on Docker Desktop & Docker CE |
| NVIDIA Container Toolkit | GPU tags only — install guide |
| NVIDIA driver ≥ 525 | Required for CUDA 12.6 compatibility |
Environment Variables (pre-set in YOLO variants)
| Variable | Value | Description |
|---|---|---|
YOLO_CONFIG_DIR | /app/weights | Ultralytics config & weights directory |
YOLO_OFFLINE | 1 | Disables Ultralytics auto-download at runtime |
License
- Image build scripts: MIT License
- PyTorch: BSD-style License
- Ultralytics YOLO (weights in YOLO variants): AGPL-3.0 License
Note: The YOLO weight files bundled in the YOLO-variant images are subject to the Ultralytics AGPL-3.0 license. Ensure your use case complies with this license.
Tag summary
Content type
Image
Digest
sha256:189308b6a…
Size
490.2 MB
Last updated
4 months ago
docker pull pytorchlab/pytorch:cpu-python3.11-yolo26m-seg