How to: GPU-Accelerated R/Keras on NixOS (via Docker)

I spent a weekend trying to get R, Keras, and CUDA to play nice together on NixOS. I suspect this’ll be the task I’m assigned on my first day in Hell.

I’ve given up on managing it natively with a shell.nix and am now using a Docker image combined with nVidia’s docker wrapper.

Before now, I’d dismissed Docker as “probable fad, soon irrelevant,”¹ so this is my first experience with it.

Enabling Docker with the nvidia wrapper took adding this to my configuration .nix:

virtualisation.docker.enable = true;
virtualisation.docker.enableNvidia = true;
hardware.opengl.driSupport32Bit = true;

I ended up cycling through a few different Docker images before rocker/ml-gpu finally worked.

sudo nvidia-docker run --gpus all --rm -v ~/:/home/rps -ti rocker/ml-gpu /bin/bash

The bit here, ~-v /:/home/rps, mounts my home directory into the Docker container, making it the least painful Linux file-sharing experience I’ve ever had, excepting maybe sshfs. I’m impressed; maybe this containerization thing isn’t so bad. This feature allowed me to easily access my training data / emacs config from inside the container.

Fixing Pillow PNG errors

Unfortunately even the Dockerized setup isn’t quite right out the box; it ships an older version of Pillow that produces errors like StopIteration: 'PngStream' object has no attribute 'chunk_bBPn'.

To fix, just install an up-to-date version into the Docker container using:

pip --upgrade pillow

How to check if Keras & R are using the GPU

Your environment ought to be working now. You can confirm that R/Keras is picking up your GPU via

library(keras)

k=backend()
sess = k$get_session()
sess$list_devices()

You can also use nvidia-smi to check the GPU’s load when you’re, say, training a model. It’s sorta like top.