Minimagen training

This repo contains the files needed to build a Docker image, which can train Minimagen on LAION-Aesthetics V2.

It does so by using the Direct Data Access (DDA) to stream the large dataset to the Docker container at runtime.

This allows you to:

• Avoid having to manually copy data to a training server, and
• Allow you to train on a subset of the data, by letting you programmatically decide which data points to use

The training script will log training parameters, metrics, and artifacts to an MLflow server connected to a DagsHub repo of your choice.

Build the Docker image

To build the Docker image as is, run:

docker build -t minimagen .

Start training

To start training, spin up a Docker container and pass in the appropriate environment variable and training parameters:

docker run --rm --name minimagen-test \
        --gpus all \
        -e DAGSHUB_TOKEN=<token> \
        -e DAGSHUB_USER_NAME=<username> \
        -e DAGSHUB_REPO_NAME=<repo_to_write_to> \
        minimagen:latest --BATCH_SIZE 2 --TIMESTEPS 25 --TESTING

Here:

• DAGSHUB_TOKEN is your access token to DagsHub (generated here)
• DAGSHUB_USER_NAME is the user name for the DagsHub account you want to use
• DAGSHUB_REPO_NAME is the repo name to send MLflow training parameters, metrics, and artifacts to
• BATCH_SIZE is the batch size to run with
• TIMESTEPS is the number of steps to train on
• TESTING limits the number of images to train/validate on to 16 (to test the pipeline)