DEV Community

Cover image for Building a PyTorch Container Image
Vultr profile image Sanskriti Harmukh
Sanskriti Harmukh for Vultr

Posted on with Aashish Chaurasiya Originally published at docs.vultr.com

Building a PyTorch Container Image

#pytorch #docker #gpu #machinelearning

PyTorch is a flexible machine learning library for building deep learning models that can perform a wide range of tasks such as image recognition and language processing. This guide walks through cloning the PyTorch source repository, building a container image for either CPU or GPU systems, and pushing the finished image to a container registry. By the end, you'll have a custom PyTorch container image built from source and published to a registry, ready to pull into any Docker-compatible environment.

Prerequisites

  • A server with Docker installed (a GPU-enabled server with the NVIDIA Container Toolkit if you're building the GPU image).
  • Access to the server using SSH as a non-root user with sudo privileges.

  • Start the Docker service.

    $ sudo systemctl start docker

  • Add the non-root user to the Docker group. For example, linuxuser.

    $ sudo usermod -aG docker linuxuser

  • Switch to the user:

    $ su - linuxuser

1. Set Up the Server

1. Create a new directory to store your PyTorch project files:

$ mkdir llama-project
Enter fullscreen mode Exit fullscreen mode

2. Switch to the directory:

$ cd llama-project
Enter fullscreen mode Exit fullscreen mode

3. Clone the PyTorch project repository using Git:

$ git clone https://github.com/pytorch/pytorch
Enter fullscreen mode Exit fullscreen mode

4. List all files and verify that a new pytorch directory is available:

$ ls
Enter fullscreen mode Exit fullscreen mode

5. Switch to the pytorch project directory:

$ cd pytorch
Enter fullscreen mode Exit fullscreen mode

6. List all directory files and verify the available Dockerfile resources:

$ ls -a
Enter fullscreen mode Exit fullscreen mode

The PyTorch project directory includes the following Dockerfile resources:

  • Dockerfile: Contains the PyTorch build context for CPU systems.
  • docker.makefile: Contains advanced GPU build context and image automation directives with support for custom dependency packages.

Use these resources in the next sections to build a CPU or GPU system container image.

2. Build a PyTorch Container Image for CPU Systems

Follow the steps below to build a new PyTorch container image for CPU systems using the Dockerfile resource on your server.

1. Build a new PyTorch container image using the Dockerfile resource. Replace pytorch-image with your desired image name:

$ docker build -f Dockerfile -t pytorch-image .
Enter fullscreen mode Exit fullscreen mode

2. View all Docker images and verify that your new PyTorch image is available:

$ docker images
Enter fullscreen mode Exit fullscreen mode

Output:

REPOSITORY     TAG                                                   IMAGE ID       CREATED       SIZE
pytorch-image  latest                                                d3674f1db4d5   2 hours ago   11.9GB
Enter fullscreen mode Exit fullscreen mode

3. Build a PyTorch Container Image for GPU Systems

Follow the steps below to build a new PyTorch container image for GPU systems using docker.Makefile, which contains the CUDA dependency build context compatible with NVIDIA GPU servers.

1. Build a new PyTorch container image for GPU using the docker.Makefile:

$ docker build -f Dockerfile -t pytorch-image .
Enter fullscreen mode Exit fullscreen mode

Wait at least 2 hours for the build process to complete. When successful, your output should look like the one below:

#29 29224.1 copying torch/include/ATen/ops/_sparse_coo_tensor_with_dims_ops.h -> build/lib.linux-x86_64-cpython-310/torch/include/ATen/ops
#29 29224.1 copying torch/include/ATen/ops/absolute_compositeimplicitautograd_dispatch.h -> build/lib.linux-x86_64-cpython-310/torch/include/ATen/ops
#29 29224.1 copying torch/include/ATen/ops/gradient_ops.h -> build/lib.linux-x86_64-cpython-310/torch/include/ATen/ops
#29 29224.1 copying torch/include/ATen/ops/crow_indices_compositeexplicitautograd_dispatch.h -> build/lib.linux-x86_64-cpython-310/torch/include/ATen/ops
#29 29224.1 copying torch/include/ATen/ops/isinf.h -> build/lib.linux-x86_64-cpython-310/torch/include/ATen/ops
#29 29224.1 copying torch/include/ATen/ops/_fore
#29 29224.1 [output clipped, log limit 2MiB reached]
#29 29235.3 warning: no previously-included files matching '*.o' found anywhere in distribution
#29 29235.4 warning: no previously-included files matching '*.dylib' found anywhere in distribution
#29 29235.7 warning: no previously-included files matching '*.swp' found anywhere in distribution
#29 DONE 29239.6s

#30 [dev 1/1] COPY --from=build /opt/conda /opt/conda
#30 DONE 10.8s

#31 exporting to image
#31 exporting layers
#31 exporting layers 41.0s done
#31 writing image sha256:d3674f1db4d5e9585a5e5b64a0cdfcb0c7d312358e267000e23b6fa9d8e7e917 done
#31 naming to docker.io/root/pytorch:v1.8.0-rc1-38416-g4dc160864b3-cuda12.1-cudnn8-devel done
#31 DONE 41.0s
Enter fullscreen mode Exit fullscreen mode

2. View all Docker images on the server and verify that a new root/pytorch image is available:

$ docker images
Enter fullscreen mode Exit fullscreen mode

Output:

REPOSITORY     TAG                                                   IMAGE ID       CREATED       SIZE
root/pytorch   v1.8.0-rc1-38416-g4dc160864b3-cuda12.1-cudnn8-devel   d3674f1db4d5   3 hours ago   20.9GB
Enter fullscreen mode Exit fullscreen mode

3. Verify the PyTorch GPU image's target CUDA version and confirm your system uses that driver version or higher:

$ cat docker.Makefile | grep CUDA_VERSION
Enter fullscreen mode Exit fullscreen mode

Output:

CUDA_VERSION_SHORT       ?= 12.1
CUDA_VERSION             ?= 12.1.1
BASE_DEVEL                = nvidia/cuda:$(CUDA_VERSION)-cudnn$(CUDNN_VERSION)-devel-ubuntu22.04
Enter fullscreen mode Exit fullscreen mode

Based on the above output, NVIDIA GPU servers with CUDA version 12.1 or higher can run your PyTorch GPU container image.

4. Upload the PyTorch Container Image to a Container Registry

1. Open your container registry's management panel and copy the registry URL, username, and API key (or password) to use when accessing the registry.

2. Switch to your server terminal session and log in to your container registry. Replace your-registry-url, exampleregistry, exampleuser, and registry-password with your actual registry details:

$ docker login https://your-registry-url/exampleregistry -u exampleuser -p registry-password
Enter fullscreen mode Exit fullscreen mode

3. Tag your PyTorch container image with your desired registry tag. For example, tag your PyTorch GPU image ID with a registry repository tag such as your-registry-url/exampleregistry/pytorch-gpu-image:

$ docker tag d3674f1db4d5 your-registry-url/exampleregistry/pytorch-gpu-image
Enter fullscreen mode Exit fullscreen mode

The above command tags the Docker image with container ID d3674f1db4d5 with your registry tag.

4. View all Docker images on the server and verify that the new tagged image is available:

$ docker images
Enter fullscreen mode Exit fullscreen mode

Output:

REPOSITORY           TAG                                                   IMAGE ID       CREATED       SIZE
root/pytorch         v1.8.0-rc1-38416-g4dc160864b3-cuda12.1-cudnn8-devel   d3674f1db4d5   4 hours ago   20.9GB
your-registry-url/exampleregistry/pytorch-gpu-image   latest                 d3674f1db4d5   4 hours ago   20.9GB
Enter fullscreen mode Exit fullscreen mode

5. Push the tagged image to your container registry repository:

$ docker push your-registry-url/exampleregistry/pytorch-gpu-image
Enter fullscreen mode Exit fullscreen mode

6. Open your container registry's management panel and check the repositories list to verify that the new repository is available.

Next Steps

  • Automate the build and push steps in a CI/CD pipeline so new PyTorch images are rebuilt on every commit.
  • Pull the image into a Kubernetes cluster or Docker Compose stack to run inference or training workloads.
  • Trim the final image size with multi-stage builds if you only need the runtime (not the full build toolchain).
  • Pin the CUDA and cuDNN versions in your Dockerfile to keep builds reproducible across environments.

For the full guide with additional tips, visit the original article on Vultr Docs.

Top comments (0)