pytorchlab/pytorch

pytorchlab/pytorch

用于GPU和CPU工作负载的最小化、生产就绪型PyTorch基础镜像——可选预捆绑YOLO权重以支持离线推理。

概述

pytorchlab/pytorch 基于官方 https://hub.docker.com/r/pytorch/pytorch 基础镜像提供精简的多变体Docker镜像。每个镜像仅添加现实世界部署所需的组件：

• 系统库：OpenCV所需的系统库（libGL、libglib2.0、libpq、libgeos）
• 清理层：构建时移除__pycache__以最小化镜像大小
• YOLO变体：预下载Ultralytics权重捆绑在镜像中，使ECS/Kubernetes工作节点在运行时无需出站依赖

可用标签

GPU（CUDA 12.6 + cuDNN 9）

基于 pytorch/pytorch:2.9.1-cuda12.6-cudnn9-runtime 构建。需要与CUDA 12.6+兼容的NVIDIA驱动。

CPU（Python 3.11 Slim）

基于 python:3.11-slim 构建，仅包含CPU版本PyTorch（比CUDA变体轻约1.5 GB）。

快速开始

bash
# GPU — 验证CUDA是否可用
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 — 轻量级推理
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__)"

作为基础镜像使用

GPU工作负载

dockerfile
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工作负载/轻量级推理

dockerfile
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"]

离线YOLO推理（运行时无需网络）

dockerfile
FROM pytorchlab/pytorch:2.9.1-cuda12.6-cudnn9-yolo26m

# 权重已预下载至/app/weights/yolo26m.pt和yolo26n.pt
# 已预设环境变量YOLO_OFFLINE=1和YOLO_CONFIG_DIR=/app/weights
WORKDIR /app
COPY . .
CMD ["python", "detect.py"]

要求

环境变量（YOLO变体中已预设）

许可证

• 镜像构建脚本：MIT许可证
• PyTorch：https://github.com/pytorch/pytorch/blob/main/LICENSE
• Ultralytics YOLO（YOLO变体中的权重）：https://github.com/ultralytics/ultralytics/blob/main/LICENSE

注意：YOLO变体镜像中捆绑的YOLO权重文件受Ultralytics AGPL-3.0许可证约束。请确保您的使用场景符合该许可证要求。