Training stylegan2 on TPU

[reidsanders]

Running eg, stylegan tests on TPU vm yields:
============================================== short test summary info ==============================================
FAILED tests/stylegan2/discriminator/test_discriminator.py::test_reference_output_afhqcat - assert DeviceArray(0.0...
FAILED tests/stylegan2/discriminator/test_discriminator.py::test_reference_output_afhqwild - assert DeviceArray(0....
FAILED tests/stylegan2/discriminator/test_discriminator.py::test_reference_output_car - assert DeviceArray(0.00418...
FAILED tests/stylegan2/discriminator/test_discriminator.py::test_reference_output_cat - assert DeviceArray(0.00065...
FAILED tests/stylegan2/discriminator/test_discriminator.py::test_reference_output_church - assert DeviceArray(0.00...
FAILED tests/stylegan2/discriminator/test_discriminator.py::test_reference_output_ffhq - assert DeviceArray(0.0080...
FAILED tests/stylegan2/discriminator/test_discriminator.py::test_reference_output_horse - assert DeviceArray(0.001...
FAILED tests/stylegan2/generator/test_generator.py::test_reference_output_afhqdog - assert DeviceArray(0.00501162,...
FAILED tests/stylegan2/generator/test_generator.py::test_reference_output_afhqwild - assert DeviceArray(0.00320451...
FAILED tests/stylegan2/generator/test_generator.py::test_reference_output_car - assert DeviceArray(0.0042854, dtyp...
FAILED tests/stylegan2/generator/test_generator.py::test_reference_output_cat - assert DeviceArray(0.00249642, dty...
FAILED tests/stylegan2/generator/test_generator.py::test_reference_output_church - assert DeviceArray(0.00342116, ...
FAILED tests/stylegan2/generator/test_generator.py::test_reference_output_cifar10 - assert DeviceArray(0.00351433,...
FAILED tests/stylegan2/generator/test_generator.py::test_reference_output_ffhq - assert DeviceArray(0.0016813, dty...
FAILED tests/stylegan2/generator/test_generator.py::test_reference_output_horse - assert DeviceArray(0.00304833, d...
FAILED tests/stylegan2/generator/test_generator.py::test_reference_output_metfaces - assert DeviceArray(0.00100499...
=============================== 16 failed, 6 passed, 2 warnings in 105.92s (0:01:45) ================================
Interestingly the tpu inference example notebook seems to work fine.

Attempting to train with

$ ~/flaxmodels/training/stylegan2$ python main.py --data_dir ~/data/data_tfrecords/ --resolution 64 --batch_size 8 --learning_rate 0.0025 --r1_gamma 0.5 --mbstd_group_size 8 --fmap_base 32768 --wandb

Gives a nan/inf error. I'm trying 256, to see if that works.

I've run this with both requirements.txt and most recent versions of tensorflow.

Is there a different setting required for using TPU? What could be causing the difference?
Thanks!

[matthias-wright]

Hi @reidsanders! I only tested it for GPUs, not for TPUs, so I only provided instructions for GPUs. It is interesting that inference works on the TPU but not training. Could you also post your Colab for training, then I can take a look. Cheers!

[reidsanders]

Well curiously I found that training stylegan2 on the tpu vm seems to work fine with 256x256 (unlike 64x64), despite the tests failing. Running tests on colab with tpu passed. Likely a library incompatibility because I was running with python 3.9 and more up to date libraries on the vm, but I haven't had a chance to track down the issue. Here's a very simple running tests in colab example
If I figure out where the failing tests with is coming from I'll post it.

Thanks!