pytorch-fips

A minimal, wolfi-based image for pytorch, a Python package that provides two high-level features: Tensor computation with strong GPU acceleration and Deep neural networks built on a tape-based autograd system.

Download this Image

The image is available on cgr.dev:

docker pull cgr.dev/chainguard/pytorch:latest

Running pytorch

PyTorch has some prerequisites which need to be configured in the environment prior to running with GPUs. For examples, please refer to TESTING.md.

Additionally, please refer to the upstream documentation for more information on configuring and using PyTorch.

Assuming the environment prerequisites have been met, below demonstrates how to launch the container:

docker run --rm -i -t \
    --privileged \
    --gpus all \
    cgr.dev/chainguard/pytorch:latest

Testing GPU Access

If your environment has connected GPUs, you can check that PyTorch has access with the following:

docker run --rm -it --gpus all cgr.dev/chainguard/pytorch:latest
Python 3.11.9 (main, Apr  2 2024, 15:40:32) [GCC 13.2.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> torch.cuda.is_available()
True
>>> torch.cuda.device_count()
1
>>> torch.cuda.get_device_name(0)
'Tesla V100-SXM2-16GB'

Testing PyTorch

As a quick intro, we will use PyTorch to create a very simple deep learning model with two linear layers and an activation function. We’ll create an instance of it and ask it to report on its parameters. Running the below will fetch a model_builder.py script from the Chainguard Images repository, place it in a folder on your host machine, and run the script in a pytorch container from a volume.

mkdir pytorch-test &&\
 curl https://raw.githubusercontent.com/chainguard-images/images/main/images/pytorch/model_builder.py > pytorch-test/model_builder.py &&\
 docker run --rm -it -v "$PWD/pytorch-test:/tmp/pytorch-test" --gpus all cgr.dev/chainguard/pytorch:latest -c "python /tmp/pytorch-test/model_builder.py"

You may also consider running this quickstart script based on the official PyTorch quickstart tutorial using the same approach as above.

Using Helm charts

As a place to get started, you may also use this Helm chart to get PyTorch running

  helm install pytorch \
  --namespace pytorch-space --create-namespace  \
  --set image.registry="cgr.dev" \
  --set image.repository="chainguard/pytorch" \
  --set image.tag=latest \
  --set containerSecurityContext.runAsUser=0 \
  --set containerSecurityContext.runAsNonRoot=false \
  --set containerSecurityContext.allowPrivilegeEscalation=true \
  --wait oci://registry-1.docker.io/bitnamicharts/pytorch