Hace unas semanas empecé a experimentar con cargas de trabajo procesadas por GPUs y me encontré un proyecto en Github que te permite colorear y restaurar audio y video.

Tengo unas fotos familiares que me gustaría restaurar; en este post voy a documentar el proceso de configurar una instancia de AWS e instalar DeOldify para restaurarlas.

Decidí utilizar AWS porque estoy esperando que todo el proceso tome un par de horas y el costo por instancia g4dn.xlarge es de $0.63 USD por hora. Tengo un equipo con GPU pero para este experimento prefiero pagar los gastos de AWS que pasar un rato configurándolo.

Voy a usar una instancia g4dn.xlarge que ofrece 4 vCPUs de un procesador Intel Xeon P-8259L a 2.5GHz, 16GB de memoria RAM y un GPU NVIDIA T4. Por simplicidad voy a usar Ubuntu 18.04 LTS.

Instalando los drivers de NVIDIA

Decidí usar la imagen de Ubuntu 18.04 que no tiene los drivers de NVIDIA instalados por defecto para documentar los pasos de instalación, así cuando lo instale en mi equipo de pruebas ya tendré una guía para seguir. Para esto estoy siguiendo los pasos descritos en el sitio oficial de NVIDIA.

Nos podríamos saltar los pasos de esta sección si usamos una de sus imágenes de Deep Learning de AWS al crear la máquina virtual.

Instalamos las cabeceras del kernel y los paquetes de desarrollo.

$ sudo apt-get install linux-headers-$(uname -r)

Nos aseguramos que los paquetes de la red de CUDA tienen prioridad sobre el repositorio de Canonical e instalamos la llave GPG de repositorio.

$ distribution=$(. /etc/os-release;echo $ID$VERSION_ID | sed -e 's/\.//g')
$ wget https://developer.download.nvidia.com/compute/cuda/repos/$distribution/x86_64/cuda-$distribution.pin
$ sudo mv cuda-$distribution.pin /etc/apt/preferences.d/cuda-repository-pin-600

Configuramos el repositorio de CUDA.

$ sudo apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/$distribution/x86_64/7fa2af80.pub
$ echo "deb http://developer.download.nvidia.com/compute/cuda/repos/$distribution/x86_64 /" | sudo tee /etc/apt/sources.list.d/cuda.list

Instalamos los drivers de CUDA.

$ sudo apt-get update
$ sudo apt-get -y install cuda-drivers

Reiniciamos y validamos que el driver esté cargado.

$ sudo reboot
# Después de reiniciar
$ cat /proc/driver/nvidia/version

# La salida se verá algo así
NVRM version: NVIDIA UNIX x86_64 Kernel Module  450.36.06  Mon Jun  1 23:19:54 UTC 2020
GCC version:  gcc version 7.5.0 (Ubuntu 7.5.0-3ubuntu1~18.04)

Instalando docker y nvidia-docker

El proyecto de DeOldify corre en Docker y requiere de nvidia-docker que es un toolkit para usar y correr contenedores Docker acelerados por GPU.

Instalamos docker.

$ sudo apt install docker.io

# Habilitamos el servicio y lo configuramos para que corra en el arranque de sistema
$ sudo systemctl start docker
$ sudo systemctl enable docker

Instalamos nvidia-docker.

# Instalamos el repositorio para nvidia-docker
$ curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | \
  sudo apt-key add -
$ distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
$ curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | \ 
  sudo tee /etc/apt/sources.list.d/nvidia-docker.list
$ sudo apt-get update

# Instalamos el paquete
$ sudo apt-get install nvidia-docker2
# Matamos el proceso actual de docker para forzar a que se reinicie
$ sudo pkill -SIGHUP dockerd

Instalar DeOldify

Ahora que tenemos las dependencias necesarias podemos pasar a instalar DeOldify con las instrucciones de su repositorio.

# Clonamos el repositorio
$ git clone https://github.com/jantic/DeOldify.git DeOldify

# Hacemos build con docker (este paso genera una imagen de ~8GB)
$ cd DeOldify && sudo docker build -t deoldify_api -f Dockerfile-api .

# Inicamos el contenedor para exponer un endpoint
$ echo "http://$(curl ifconfig.io):5000" && nvidia-docker run --ipc=host -p 5000:5000 -d deoldify_api

NOTA: Al correr el contenedor en modo API no hace uso del GPU por defecto, hay que agregar las siguientes líneas en los imports de `api.py` para que haga uso del GPU.

from deoldify import device
from deoldify.device_id import DeviceId
device.set(device=DeviceId.GPU0)

¡A colorear fotos!

Ya que todo está corriendo como lo esperamos, podemos usar la URL regresada por el último comando de la instalación de DeOldify para procesar las fotos.

Desde la terminal puedo llamar con curl la API que expone el contenedor de DeOldify.

$ curl -X POST "http:/{IP_SERVIDOR}:5000/process" -H "accept: image/png" -H "Content-Type: application/json" -d "{\"source_url\":\"https://{IMAGEN_A_PROCESAR}\", \"render_factor\":35}" --output output.png 

Convertí 29 fotos y el tiempo de respuesta fue entre 1 y 7 segundos por foto (entre 10 y 38 segundos en CPU). En algunos casos tuve que ajustar el valor de `render_factor` para mejorar los resultados; conforme se incrementa este valor el uso de memoria y tiempo de procesamiento aumenta.

Original

Coloreada

Coloreando videos

Es posible también utilizar el mismo paquete para colorear videos, en este caso, aunque no es requerido, es importante utilizar el GPU, de lo contrario tomaría bastante tiempo en completar.

Es necesario hacer los siguientes cambios menores al paquete y generar una nueva imagen para docker (ambos contenedores pueden estar corriendo al mismo tiempo).

Primero vamos a agregar las siguientes cabeceras al archivo api-video.py.

from deoldify import device
from deoldify.device_id import DeviceId
device.set(device=DeviceId.GPU0)

Hacemos una copia del archivo Dockerfile-api en Dockerfile-apivideo y cambiamos la última línea de CMD ["app.py"] a CMD ["app-video.py"].

Después dentro del directorio DeOldify corremos el siguiente comando para generar la imagen nueva.

$ sudo docker build -t deoldify_api_video -f Dockerfile-apivideo .

Y una vez que termine iniciamos otro contenedor, esta vez utilizando el puerto 5001 en caso de que queramos tener ambos corriendo al mismo tiempo.

$ echo "http://$(curl ifconfig.io):5001" && nvidia-docker run --ipc=host -p 5001:5000 -d deoldify_api_video

Ahora sí, podemos llamar la API.

$ curl -X POST "http:/{IP_SERVIDOR}:5001/process" -H "accept: application/octet-stream" -H "Content-Type: application/json" -d "{\"source_url\":\"https://{VIDEO_A_PROCESAR}\", \"render_factor\":35}" --output output.mp4 

Hice una prueba con un video que publicaron en Reddit de 10:09 minutos de duración y tomó 1:41 horas en procesar. Si hubiera decidido procesar esto con el CPU, según mis cuentas de servilleta, hubiera tomado 76:13 horas. Este fue el resultado:

 

 

Los resultados de las fotos y video son mucho mejor de lo que esperaba y una vez configurado el contenedor no es necesario hacer ningún ajuste. Voy a instalarlo en mi equipo de pruebas local para implementar un par de mejoras a la API, y probablemente luego me aventure a entrenar un modelo con otros videos.