StyleGAN2-ADA only work with Tensorflow 1. Run the next cell before anything else to make sure we’re using TF1 and not TF2.
%tensorflow_version 1.x
!nvidia-smi
Colab is a little funky with training. I’ve found the best way to do this is to install the repo directly into your Google Drive folder.
First, mount your Drive to the Colab notebook:
from google.colab import drive
drive.mount('/content/drive')
Next, run this cell. If you’re already installed the repo, it will skip the installation process and change into the repo’s directory. If you haven’t installed it, it will install all the files necessary.
import os
if os.path.isdir("/content/drive/My Drive/colab-sg2-ada"):
%cd "/content/drive/My Drive/colab-sg2-ada/stylegan2-ada"
else:
#install script
%cd "/content/drive/My Drive/"
!mkdir colab-sg2-ada
%cd colab-sg2-ada
!git clone https://github.com/dvschultz/stylegan2-ada
%cd stylegan2-ada
!mkdir downloads
!mkdir datasets
%cd "/content/drive/My Drive/colab-sg2-ada/stylegan2-ada"
!git config --global user.name "test"
!git config --global user.email "test@test.com"
!git fetch origin
!git checkout origin/main -- train.py
Note: You only need to do this once per dataset. If you have already run this and are returning to conntinue training, skip these cells.
Next we need to convert our image dataset to a format that StyleGAN2-ADA can read from. There are two options here. You can upload your dataset directly to Colab (as a zipped file), or you can upload it to Drive directly and read it from there.
#if you manually uploaded your dataset to Colab, unzip it
zip_path = "/content/CAT1024.zip"
!unzip {zip_path} -d /content/
Now that your image dataset is uploaded, we need to convert it to the .tfrecords
format.
Depending on the resolution of your images and how many you have, this can take a while.
#update this to the path to your image folder
dataset_path = "/content/CAT1024"
#give your dataset a name
dataset_name = "cat1024"
#you don't need to edit anything here
!python dataset_tool.py create_from_images ./datasets/{dataset_name} {dataset_path}
We’re ready to start training! There are numerous arguments to training, what’s listed below are the most popular options. To see all the options, run the following cell.
!python train.py --help
#this name must EXACTLY match the dataset name you used when creating the .tfrecords file
dataset_name = "dante1024"
#how often should the model generate samples and a .pkl file
snapshot_count = 4
#should the images be mirrored left to right?
mirrored = True
#should the images be mirrored top to bottom?
mirroredY = True
#metrics?
metric_list = None
#augments
augs = "bg"
#
# this is the most important cell to update
#
# running it for the first time? set it to ffhq(+resolution)
# resuming? get the path to your latest .pkl file and use that
resume_from = "/content/drive/My\ Drive/colab-sg2-ada2/stylegan2-ada/results/00008-dante1024-mirror-mirrory-11gb-gpu-bg-resumecustom/network-snapshot-000160.pkl"
#don't edit this unless you know what you're doing :)
!python train.py --outdir ./results --snap={snapshot_count} --cfg=11gb-gpu --data=./datasets/{dataset_name} --augpipe={augs} --mirror={mirrored} --mirrory={mirroredY} --metrics={metric_list} --resume={resume_from} --augpipe="bg"
Once the above cell is running you should be training!
Don’t close this tab! Colab needs to be open and running in order to continue training. Every ~15min or so a new line should get added to your output, indicated its still training. Depending on you snapshot_count
setting you should see the results folder in your Google drive folder fill with both samples (fakesXXXXXx.jpg
) and model weights (network-snapshot-XXXXXX.pkl
). The samples are worth looking at while it trains but don’t get too worried about each individual sample.
If you chose a metric, you will also see scores for each snapshot. Don’t obsess over these! they are a guide, it can go up or down slightly for each snapshot. What you want to see is a gradual lowering of the score over time.
Once Colab shuts off, you can Reconnect the notebook and re-run every cell from top to bottom. Make sure you update the resume_from
path to continue training from the latest model.