Skip to content

Set up a GPU data science project with pixi

by Tim Hopper · Markdown
Scientific Python verified with pixi 0.83

This tutorial builds a PyTorch image classification project managed by pixi. It loads a pretrained ResNet model, classifies a sample image, and prints the top-5 predictions. The project uses pixi’s multi-environment feature to support both GPU (CUDA) and CPU-only machines from a single pixi.toml.

Prerequisites

Why conda-forge for PyTorch

PyTorch on conda-forge shares the CUDA runtime with other packages. The solver selects compatible CUDA runtime libraries alongside PyTorch and cuDNN. A complete development toolkit is unnecessary for running prebuilt packages; the host still needs an NVIDIA driver. A solved environment does not guarantee driver or hardware compatibility.

PyPIThe Python Package Index, the public repository where Python packages are published and downloaded from. "pip install requests" fetches requests from PyPI. Learn more → PyTorch wheelsA prebuilt Python package file (.whl) that installs without compiling anything. The standard distribution format for Python packages. Learn more → also supply CUDA libraries, bundled or through Python package dependencies. See How to Install PyTorch with uv for the PyPI approach.

Create the project 

$ pixi init image_classifier
✔ Created /path/to/image_classifier/pixi.toml
$ cd image_classifier

If you see pixi: command not found, restart your shell so the installer’s PATH update takes effect, then try again.

Notice the .gitignore and pixi.toml files pixi created. The .gitignore is pre-configured to exclude the .pixi/ environment directory from version control.

Define CPU and GPU environments

Replace pixi.toml with this complete manifest:

pixi.toml
[workspace]
name = "image_classifier"
channels = ["conda-forge"]
platforms = [
  { name = "linux-cpu", platform = "linux-64" },
  { name = "linux-cuda", platform = "linux-64", cuda = "12.8" },
]

[dependencies]
python = "3.12.*"
torchvision = "*"
pillow = "*"

[feature.cpu]
platforms = ["linux-cpu"]

[feature.cpu.dependencies]
pytorch-cpu = "*"

[feature.cuda]
platforms = ["linux-cuda"]

[feature.cuda.dependencies]
pytorch-gpu = "*"
cuda-version = "12.8.*"

[environments]
default = ["cpu"]
gpu = ["cuda"]

pytorch-cpu and pytorch-gpu select the backend explicitly. cuda-version constrains CUDA packages; it does not itself force GPU PyTorch. Both environments inherit the shared dependenciesAn external package your project needs, listed in pyproject.toml so tools can install it automatically. and solve independently so incompatible CPU/GPU builds do not share one solve group.

The named CUDA platform declares a driver capability for solving and host validation. It does not install a driver or prove a GPU operation works. cuda-version constrains the environment’s CUDA packages; torch.version.cuda can report an older compatible build, so inspect it rather than assuming the constraint equals PyTorch’s compile-time version.

pixi install
pixi run python -c "import torch; assert torch.version.cuda is None; print('CPU build selected')"

Expect CPU build selected. Commit the generated pixi.lock to record the compatible versions selected from conda-forge. Python 3.12 avoids relying on immediate wheel/package availability for the newest Python release.

On the NVIDIA host, install and check the GPU environment:

pixi install --environment gpu
pixi run --environment gpu python -c "import torch; assert torch.cuda.is_available(); print(torch.__version__, torch.version.cuda, torch.cuda.get_device_name(0))"

Expect a CUDA-enabled torch version, its runtime version, and the GPU name. If Pixi rejects the system requirements, check the host driver. For a GPU-less image build, CONDA_OVERRIDE_CUDA=12.8 can declare the target capability during installation; remove the override on the execution host. An override bypasses detection and is not evidence of GPU availability.

Download a sample image

Create a data directory and download a test image:

mkdir data
curl -L -o data/cat.jpg "https://upload.wikimedia.org/wikipedia/commons/thumb/4/4d/Cat_November_2010-1a.jpg/1200px-Cat_November_2010-1a.jpg"

Write the classification script

Create classify.py:

classify.py
import argparse
import json
import torch
import torchvision.models as models
import torchvision.transforms as transforms
from PIL import Image


def load_labels():
    """Load ImageNet class labels."""
    url = "https://raw.githubusercontent.com/anishathalye/imagenet-simple-labels/master/imagenet-simple-labels.json"
    import urllib.request
    with urllib.request.urlopen(url) as response:
        return json.loads(response.read().decode())


def classify_image(image_path, backend):
    if backend == "cuda":
        assert torch.cuda.is_available(), "GPU environment cannot access CUDA"
    device = torch.device(backend)
    print(f"Using device: {device}")

    # Load pretrained ResNet
    model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
    model.to(device)
    model.eval()

    # Preprocess image
    preprocess = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225],
        ),
    ])

    image = Image.open(image_path).convert("RGB")
    input_tensor = preprocess(image).unsqueeze(0).to(device)

    # Run inference
    with torch.no_grad():
        output = model(input_tensor)
    assert output.device.type == backend
    assert torch.isfinite(output).all()

    probabilities = torch.nn.functional.softmax(output[0], dim=0)
    top5_prob, top5_idx = torch.topk(probabilities, 5)

    labels = load_labels()

    print(f"\nTop 5 predictions for {image_path}:")
    for i in range(5):
        idx = top5_idx[i].item()
        prob = top5_prob[i].item()
        print(f"  {labels[idx]:30s} {prob:.1%}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("image_path")
    parser.add_argument("--device", choices=["cpu", "cuda"], default="cpu")
    args = parser.parse_args()
    classify_image(args.image_path, args.device)

Run in the default (CPU) environment 

pixi run python classify.py data/cat.jpg

On first run, PyTorch downloads the pretrained ResNet-50 weights (~98 MB) before classifying:

Downloading: "https://download.pytorch.org/models/resnet50-11ad3fa6.pth" to /home/user/.cache/torch/hub/checkpoints/resnet50-11ad3fa6.pth
100%|██████████████| 97.8M/97.8M [00:08<00:00, 12.2MB/s]

The cache path depends on TORCH_HOME and XDG_CACHE_HOME; the default is ~/.cache/torch/hub/. Subsequent runs reuse the cached weights.

Expected output after the weights are cached:

Using device: cpu

Top 5 predictions for data/cat.jpg:
  tabby cat                      47.2%
  tiger cat                      30.1%
  Egyptian cat                   8.5%
  lynx                           1.2%
  Persian cat                    0.4%

The exact percentages vary slightly across PyTorch versions.

Run in the GPU environment

On a machine with an NVIDIA GPU and CUDA drivers:

pixi run --environment gpu python classify.py data/cat.jpg --device cuda

Expect Using device: cuda and five predictions. The script asserts that inference produced a finite CUDA tensor; an unavailable GPU raises instead of silently running on CPU.

Add a task for classification

pixi.toml
[tasks]
classify = "python classify.py data/cat.jpg"
pixi run classify

The output is the same as running pixi run python classify.py data/cat.jpg directly. Named tasks are aliases that give collaborators discoverable entry points.

Add Jupyter for exploration 

$ pixi add jupyter
✔ Added jupyter >=1.1.1,<2

Launch Jupyter:

pixi run jupyter lab

Jupyter prints a local URL (e.g. http://localhost:8888/lab?token=...) and opens a browser tab automatically. If no browser opens, copy the URL from the terminal and paste it. Press Ctrl+C in the terminal to stop the server.

Within a notebook, you can import torch and experiment with different models, images, or preprocessing pipelines. The notebook kernel uses the project’s environment automatically. Jupyter includes IPython as its interactive shell.

Final project structure

    • pixi.toml
    • pixi.lock
    • classify.py
    • .gitignore
      • cat.jpg

Next steps

Last updated on