A Docker Image for Computer Vision

语速

When debugging deep learning code, we often face headaches due to environment issues.

To facilitate debugging, packaging environments like PyTorch and CUDA into a Docker image is an excellent choice.

Why?

Time-saving: Repeatedly configuring and adjusting versions wastes time, leading to spending a lot of effort on ops tasks.
Environment stability: Once a Docker image is built, it is static and can be pulled directly.
Easy migration: Pre-configured environments can be migrated across different machines.

How to Build

Here is an example Docker image for packaging a deep learning environment:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
# Change to your desired pytorch version
FROM pytorch/pytorch:2.4.1-cuda11.8-cudnn9-devel

# These are commonly used packages
RUN apt-get update && apt-get install git zsh ffmpeg libsm6 libxext6  -y && apt-get clean && rm -rf /var/lib/apt/lists/*

WORKDIR /app

# Place at the root of the codebase to install requirements.txt
COPY requirements.txt .
RUN pip install -r requirements.txt

# install jupyterlab
RUN pip install jupyterlab
# COPY . .

# Use jupyterlab to host, can start quickly, token is `yourtoken`. If you use it on the public network, consider using a more complex token.
CMD ["jupyter", "lab", "--ip=0.0.0.0", "--port=8888", "--no-browser", "--allow-root", "--NotebookApp.token=yourtoken"]

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
services:
  notebook:
    build:
      context: .
      dockerfile: Dockerfile
    volumes: # You can also mount the dataset you need
      - .:/app
      - ~/.ssh:/root/.ssh # Support ssh
    ports:
      - 8888:8888
    shm_size: '32gb'
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]

This example installs some basic libraries, and opencv-python can be installed via pip.

Place the Dockerfile in the directory, then you can start it using docker compose.

The startup command is: docker compose up -d.

Download from Dockerhub

To make it convenient for everyone to use directly, I have packaged this image and uploaded it to Dockerhub. The download command is:

1
docker pull svtter/debian-pytorch

The source code can be obtained from here:

Using on Runpod

For everyone’s convenience, I have created a template on Runpod.

https://console.runpod.io/deploy?template=m0shpm3vgg&ref=g5qp1x9x

You can directly use this image by using this template.