docker-stylegan2-ada icon indicating copy to clipboard operation
docker-stylegan2-ada copied to clipboard

My Docker image for running Stylegan2 ADA with GPU

Docker file to run StyleGAN2 ADA, docker-stylegan2-ada

  • Licensed under the MIT license
  • GitHub: https://github.com/jeffheaton/docker-stylegan2-ada
  • DockerHub: https://hub.docker.com/r/heatonresearch/stylegan2-ada

Jeff Heaton's Docker image for running Stylegan2 ADA with GPU. This dockerfile is intended to be run from NVIDIA Docker. To start a container from this image run something similar to the following:

nvidia-docker run -it -u $(id -u):$(id -g) -v /home/username/:/mnt --device /dev/nvidia0:/dev/nvidia0 --device /dev/nvidiactl:/dev/nvidiactl --device /dev/nvidia-uvm:/dev/nvidia-uvm heatonresearch/stylegan2-ada /bin/bash

The above command establishes the following:

  • -u $(id -u):$(id -g) - This causes the user to be logged into Docker to be the same as the user running the Docker command. This ensures that all permissions and ownerships are correct on the mounted volumes.
  • /home/username/:/mnt - You should mount a volume to access your images and store results to.
  • /bin/bash - We start the BASH shell to allow you to execute commands.

Convert Image Data

Your images must be converted to TFRecords. It is very important that all of your images be of the same size, be square, and of the same color depth. You cannot mix color and greyscale images. The following command converts a data set named "fish". The source JPEGs should be in the /mnt/data/fish and the resulting TFRecords will be written to /mnt/datasets/fish.

cd /home/stylegan2-ada/
python dataset_tool.py create_from_images /mnt/datasets/fish /mnt/data/fish

Train GANs

To actually train your GAN use a command similar to the following.

cd /home/stylegan2-ada/
python train.py --gpus=1 --data=/mnt/datasets/fish --mirror=1 --kimg 3000 --outdir=/mnt/results --aug=ada

You must provide an input directory that contains your image TFRecords. You must also provide an output results directory. The results directory will contain snapshots of your generator, sample images as training progresses, and a log.