andresfr/mlframeworks repository overview

⁠Docker image with GPU-compatible ML Python libraries

A container that puts together PyTorch, OpenCV and Caffe on the same GPU. This can be useful by itself to combine third-party projects from different sources, but also as a baseline that can be altered to have different frameworks and versions.

Versions (check Dockerfile):

• Ubuntu 18.04
• CUDA 11.1
• CUDNN 8
• PyTorch 1.7.0
• TensorFlow 2
• OpenCV 4

The image contains the Dockerfile and also the source code for the frameworks. Approximate size on Docker Host: 19GB: 9 GB from the NVIDIA "parent" image (see here⁠) and 10 from the extra libraries. This can be improved by pruning unnecessary source code (OpenCV is very heavyweight). Suggestions are welcome!

⁠Downloading and Running:

Assuming your Docker, Nvidia driver and Nvidia-container facilities are up and running (see e.g. here⁠), you can build the latest version of this container from this repo as follows:

docker build -t andresfr/mlframeworks https://gitlab.eps.surrey.ac.uk/templates/docker_ml_frameworks.git

Alternatively, you can directly download a working snapshot (faster, but not guaranteed to be the latest) from dockerhub⁠ as follows:

docker pull andresfr/mlframeworks

You can then start an interactive session as follows:

docker run --gpus all -it --ipc=host --rm andresfr/mlframeworks python

# Try the following!
import cv2
import caffe
import tensorflow as tf
import torch

print(cv2.getBuildInformation())
caffe.set_mode_gpu()
caffe.set_device(0)
tf.test.is_gpu_available()
print("PyTorch CUDA:", torch.cuda.is_available())