Cómo entrenar un modelo de Machine Learning sin casarte con una GPU

Cuando queremos desarrollar un modelo de machine learning, usualmente experimentamos con Jupyter Notebooks, una plataforma que resulta muy cómoda para esa tarea por el feedback inmediato que nos provee. Sin embargo, en algún momento necesitaremos entrenar ese modelo por un largo rato y esto usualmente requiere que tengamos una GPU a disposición.

¿Qué opciones tenemos para resolver este problema?

Para esto contamos con varias alternativas:

1. Comprar una computadora con GPU para tener en nuestro lugar de trabajo.
2. Contratar un servicio que, por una suscripción mensual, nos dé acceso a una plataforma para correr nuestros notebooks en GPUs.
3. Utilizar algún servicio de cómputo en la nube más tradicional donde contratamos una computadora con GPU y pagamos por tiempo de uso, y a la que tenemos acceso remoto (por ejemplo, via SSH).

En nuestro caso descartamos la primera opción por ser costosa y poco práctica para desarrollar en grupo. Durante unos meses contratamos Paperspace (opción 2). Sin embargo, al tener nuestros datos guardados en un bucket de Google Cloud (GCS), nos topábamos con problemas de permisos frecuentemente. Además de eso, las GPU del plan básico de Paperspace no estaban disponibles con frecuencia y debíamos copiar manualmente los datos de GCS a la máquina de Paperspace cada vez que comenzábamos un notebook nuevo, lo cual es un proceso lento y engorroso. La opción 3 parecía la única alternativa posible, pero tiene la desventaja de que debemos elegir un tipo de recurso y quedarnos con él (GPU, CPU, RAM, etc.) a pesar de que nuestras necesidades vayan variando. Además de eso, tenemos que acordarnos de apagar la máquina una vez que el entrenamiento terminó.

¿Y si ninguna de estas opciones nos satisface?

Elegimos entonces utilizar el servicio de Google llamado Vertex AI Training, que permite entrenar modelos en la infraestructura deseada sin levantar una máquina en la nube: es un servicio serverless. Las ventajas de esta solución son dos:

1. Podemos elegir la infraestructura en la que queremos entrenar nuestro modelo cada vez, para ahorrar costos.
2. Vertex AI está integrado con los servicios de almacenamiento de Google. Por lo tanto, si nuestros datos están guardados en un bucket de GCS, podremos acceder a los mismos prácticamente sin hacer nada.

Paso 1: containerizar nuestro ambiente de entrenamiento

Supongamos que tenemos nuestros datos en un bucket llamado data-bucket, en el marco de un proyecto de Google Cloud llamado proyecto-de-ml y para correr el entrenamiento de nuestro modelo invocamos a train.py:

# Con esta línea tenemos acceso a la ruta del
# bucket para tomar y guardar datos
BUCKET_ROOT = os.getenv['BUCKET_ROOT']
# Instanciamos nuestro modelo
model = ...
# Ruta del dataset (ejemplo)
dataset_path = BUCKET_ROOT+"/data"
# Entrenamiento
model.train(data=dataset_path, epochs=100)
# Guardamos los resultados
date_timestamp = datetime.now().strftime("%d%m%Y_%H%M")
torch.save(model.state_dict(), f"{BUCKET_ROOT}/runs/{date_timestamp}")

Vamos a crear un Dockerfile que nos permite resumir todo lo necesario para entrenar nuestro modelo en pocas líneas de código. De esta manera, cuando la utilicemos en Vertex AI Training, este servicio la levantará del GAR (Google Artifact Registry, similar a Dockerhub) y tendrá todo instalado y listo para correr train.py. En nuestro caso estamos utilizando Pytorch, pero podemos encontrar imágenes que tienen otros frameworks instalados.

# Imagen base
FROM gcr.io/deeplearning-platform-release/pytorch-gpu

ENV BUCKET_ROOT=/gcs/data-bucket
WORKDIR /

# Instalamos dependencias (opcional)
RUN apt-get update && apt-get install ...
RUN pip install ...

# Copiamos train.py a nuestro container
COPY train.py /train.py

# Entrenamos el modelo
ENTRYPOINT ["python", "train.py"]

Paso 2: pushear nuestra imagen

Para poder encontrar nuestra imagen de Docker en el repositorio de GAR, primero lo creamos (yo lo llamo imagenes-entrenamiento) así:

gcloud artifacts repositories create imagenes-entrenamiento \
--repository-format=docker --location=us-central1

Configuramos Docker con gcloud auth configure-docker us-central1-docker.pkg.dev y, por último, buildeamos y pusheamos:

IMAGE_URI=us-central1-docker.pkg.dev/proyecto-de-ml/imagenes-entrenamiento/entrenamiento:latest
docker build ./ -t $IMAGE_URI
docker push $IMAGE_URI

Paso 3: entrenar el modelo

Una vez que tenemos nuestra imagen en GAR, ya estamos listos para entrenar. En la sección de Vertex AI training de la consola de Google Cloud, creamos un pipeline:

Después:

• Elegimos “No managed dataset”, seleccionamos “Custom training (advanced)” como método de entrenamiento y “Continue”.

s

• Seleccionamos “Train new model”, elegimos un nombre para nuestro modelo y “Continue”
• En el paso “Training container” seleccionamos “Custom container” y elegimos nuestra imagen de Docker. En “Model output directory” seleccionamos nuestro bucket. Dejamos todo el resto en blanco y continuamos.
• Salteamos la parte de “Hyperparameters”.
• Elegimos nuestro tipo de máquina.
• Salteamos el resto y ¡entrenamos!

Si entramos a nuestro bucket veremos los resultados del entrenamiento y no tendremos que apagar ninguna máquina.

En resumen…

Esta es una forma bastante simple de entrenar modelos de ML sin casarnos con una GPU. Si tenemos un script para esta tarea, pero solamente necesitamos los recursos de una GPU de vez en cuando (por ejemplo, si lo entrenamos durante pocas épocas en nuestra máquina local para hacer pruebas pero hacemos un entrenamiento más largo para la versión definitiva), utilizar una infraestructura serverless como la de VertexAI nos permite:

Tener todos nuestros datos y nuestra infra en la misma plataforma, lo cual nos ahorra problemas de permisos

Ahorrar costos porque solo pagamos por lo que utilizamos y elegimos nuestros recursos según la necesidad del momento.