# Building a PyTorch Container Image

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, following ML container build practices documented in [**Vultr Docs**](https://docs.vultr.com/how-to-build-a-pytorch-container-image).

---

## 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.

  ```console
  sudo systemctl start docker
  ```
- Add the non-root user to the Docker group. For example, `linuxuser`.

  ```console
  sudo usermod -aG docker linuxuser
  ```
- Switch to the user:

  ```console
  su - linuxuser
  ```

## 1\. Set Up the Server

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

```console
mkdir llama-project
```

**2\. Switch to the directory:**

```console
cd llama-project
```

**3\. Clone the PyTorch project repository using Git:**

```console
git clone https://github.com/pytorch/pytorch
```

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

```console
ls
```

**5\. Switch to the** `pytorch` **project directory:**

```console
cd pytorch
```

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

```console
ls -a
```

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:**

```console
docker build -f Dockerfile -t pytorch-image .
```

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

```console
docker images
```

Output:

```plaintext
REPOSITORY     TAG                                                   IMAGE ID       CREATED       SIZE
pytorch-image  latest                                                d3674f1db4d5   2 hours ago   11.9GB
```

## 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`**:**

```console
docker build -f Dockerfile -t pytorch-image .
```

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

```plaintext
#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
```

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

```console
docker images
```

Output:

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

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

```console
cat docker.Makefile | grep CUDA_VERSION
```

Output:

```plaintext
CUDA_VERSION_SHORT       ?= 12.1
CUDA_VERSION             ?= 12.1.1
BASE_DEVEL                = nvidia/cuda:$(CUDA_VERSION)-cudnn$(CUDNN_VERSION)-devel-ubuntu22.04
```

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:**

```console
docker login https://your-registry-url/exampleregistry -u exampleuser -p registry-password
```

**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`**:**

```console
docker tag d3674f1db4d5 your-registry-url/exampleregistry/pytorch-gpu-image
```

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:**

```console
docker images
```

Output:

```plaintext
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
```

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

```console
docker push your-registry-url/exampleregistry/pytorch-gpu-image
```

**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**](https://docs.vultr.com/how-to-build-a-pytorch-container-image).