PyTorch CPU-only installation with uv ⚡️ and Docker 🐋

Dockerfile

ARG TORCH_VERSION=2.2.1

FROM python:3.12@sha256:35eff340c0acd837b7962f77ee4b8869385dd6fe7d3928375a08f0a3bdd18beb AS basic

# Install uv
# https://github.com/astral-sh/uv
ENV VIRTUAL_ENV=/usr/local
ADD --chmod=755 https://astral.sh/uv/install.sh /install.sh
RUN /install.sh && \
	rm install.sh

# Install PT (CPU-only)
ARG TORCH_VERSION
SHELL ["/bin/bash", "-c"]
RUN source "$HOME/.cargo/env" && \
	uv pip install "torch==${TORCH_VERSION}+cpu" --no-cache -f https://download.pytorch.org/whl/torch_stable.html

# Check installation
RUN wget --progress=dot:giga https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py
CMD ["python", "collect_env.py"]

FROM python:3.12-alpine3.19@sha256:849ed6079c9f797ca9c1b7d6aea1c00aea3ac35110cbd0d6003f15950017ea8d AS alpine

# Install uv
# https://github.com/astral-sh/uv
ENV VIRTUAL_ENV=/usr/local
ADD --chmod=755 https://astral.sh/uv/install.sh /install.sh
RUN /install.sh && \
	rm install.sh

# Install PT (CPU-only)
ARG TORCH_VERSION
RUN . "$HOME/.cargo/env" && \
	uv pip install "torch==${TORCH_VERSION}+cpu" --no-cache -f https://download.pytorch.org/whl/torch_stable.html

# Check installation
RUN wget -q https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py
CMD ["python", "collect_env.py"]

Just run docker build -t test:basic --target basic . to build the image and docker run --rm test:basic to check the installation

Output:

Collecting environment information...
PyTorch version: 2.2.1+cpu
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: Debian GNU/Linux 12 (bookworm) (x86_64)
GCC version: (Debian 12.2.0-14) 12.2.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.36

Python version: 3.12.2 (main, Feb 13 2024, 09:17:46) [GCC 12.2.0] (64-bit runtime)
Python platform: Linux-5.15.133.1-microsoft-standard-WSL2-x86_64-with-glibc2.36
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                       x86_64
CPU op-mode(s):                     32-bit, 64-bit
Address sizes:                      48 bits physical, 48 bits virtual
Byte Order:                         Little Endian
CPU(s):                             2
On-line CPU(s) list:                0,1
Vendor ID:                          AuthenticAMD
Model name:                         AMD Ryzen 5 PRO 4650U with Radeon Graphics
CPU family:                         23
Model:                              96
Thread(s) per core:                 2
Core(s) per socket:                 1
Socket(s):                          1
Stepping:                           1
BogoMIPS:                           4191.99
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl tsc_reliable nonstop_tsc cpuid extd_apicid pni pclmulqdq ssse3 fma cx16 sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm cmp_legacy cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw topoext ssbd ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 clzero xsaveerptr virt_ssbd arat umip rdpid
Hypervisor vendor:                  Microsoft
Virtualization type:                full
L1d cache:                          32 KiB (1 instance)
L1i cache:                          32 KiB (1 instance)
L2 cache:                           512 KiB (1 instance)
L3 cache:                           4 MiB (1 instance)
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit:        Not affected
Vulnerability L1tf:                 Not affected
Vulnerability Mds:                  Not affected
Vulnerability Meltdown:             Not affected
Vulnerability Mmio stale data:      Not affected
Vulnerability Retbleed:             Mitigation; untrained return thunk; SMT enabled with STIBP protection
Vulnerability Spec rstack overflow: Mitigation; safe RET
Vulnerability Spec store bypass:    Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:           Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:           Mitigation; Retpolines, IBPB conditional, STIBP always-on, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Not affected

Versions of relevant libraries:
[pip3] torch==2.2.1+cpu
[conda] Could not collect