Using uv with PyTorch

uv significantly speeds up my pytorch in docker builds

  # Setup virtual env
  ENV VIRTUAL_ENV=/app/.venv
  ENV PATH="$VIRTUAL_ENV/bin:$PATH"
  RUN python3 -m venv $VIRTUAL_ENV
  RUN . $VIRTUAL_ENV/bin/activate

  # install using uv
  RUN pip install uv
  RUN uv pip install torch==${TORCH_VERSION} --index-url https://download.pytorch.org/whl/cpu

The index-url makes it really convenient.

Use —copy-from to make it even faster, and use a cache mount

https://docs.astral.sh/uv/guides/integration/docker/#install...

Speeds up installation, or speeds up PyTorch in general?

Almost certainly only the install. Uv is basically a pip tool substitute with a few other bells and whistles too but shouldn't affect run time whatsoever.

Also use pyproject.toml to specify dependencies, not manually installing stuff with uv pip

Apparently, uv respects a key in project.toml:

    [[tool.uv.index]]
    name = "pytorch-cu124"
    url = "https://download.pytorch.org/whl/cu124"
    explicit = true

Nice, I wonder if there is a way to make it conditional, e.g. to pick a different key for a cpu vs cuda build.

Eh, in that case the environment variable is a bit more evident in a Dockerfile.

That’s basically what TFA is about? Set up multiple index urls, and use a bog-standard platform marker on the PyTorch dep to pick between them.