Abstract
StyleGAN2 is a Tensorflow-based Generative Adversarial Network (GAN) framework that represents the state-of-the-art in generative image modelling. The current release of StyleGAN2 implements multi-GPU training via Tensorflow’s device contexts which limits data parallelism to a single node. In this work, a data-parallel multi-node training capability is implemented in StyleGAN2 via Horovod which enables harnessing the compute capability of larger cluster architectures. We demonstrate that the new Horovod-based communication outperforms the previous context approach on a single node. Furthermore, we demonstrate that the multi-node training does not compromise the accuracy of StyleGAN2 for a constant effective batch size. Finally, we report strong and weak scaling of the new implementation up to 64 NVIDIA Tesla A100 GPUs distributed across eight NVIDIA DGX A100 nodes, demonstrating the utility of the approach at scale.
This is a preview of subscription content, log in via an institution to check access.
Access this chapter
Tax calculation will be finalised at checkout
Purchases are for personal use only
Similar content being viewed by others
References
Goodfellow, I., et al.: Generative adversarial nets. In: Advances in Neural Information Processing Systems 27, Ghahramani, Z., Welling, M., Cortes, C., Lawrence, N.D., Weinberger, N.Q., (Eds.), Curran Associates Inc, pp. 2672–2680 (2014). http://papers.nips.cc/paper/5423-generative-adversarial-nets.pdf
Karras, T., Aila, T., Laine, S., Lehtinen, J.: Progressive growing of gans for improved quality, stability, and variation, CoRR.abs/1710.10196 (2017). http://arxiv.org/abs/1710.10196
Brock, A., Donahue, J., Simonyan, K.: Large scale GAN training for high fidelity natural image synthesis. In: Proceedings of ICLR (2019)
Karras, T., Laine, S., Aila, T.: A style-based generator architecture for generative adversarial networks. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), June 2019
Karras, T., Laine, S., Aittala, M., Hellsten, J., Lehtinen, J., Aila, T.: Analyzing and improving the image quality of stylegan. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), June 2020
Kingma, D.P., Rezende, D.J., Mohamed, S., Welling, M.: Semi-supervised learning with deep generative models, CoRR.abs/1406.5298 (2014). http://arxiv.org/abs/1406.5298
van den Oord, A., Vinyals, O., Kavukcuoglu, K.: Neural discrete representation learning. In: Proceedings of NIPS (2017)
Razavi, A., van den Oord, A., Vinyals, O.: Generating diverse high-fidelity images with vq-vae-2. In: Proceedings of NeurIPS (2019)
van den Oord, A., Kalchbrenner, N., Kavukcuoglu, K.: Pixel recurrent neural networks. In: ICML, pp. 1747–1756 (2016)
van den Oord, A., Kalchbrenner, N., Vinyals, O., Espeholt, L., Graves, A., Kavukcuoglu, K.: Conditional image generation with PixelCNN decoders, CoRR.abs/1606.05328 (2016)
Dinh, L., Sohl-Dickstein, J., Bengio, S.: Density estimation using Real NVP. CoRR. abs/1605.08803 (2016)
Kingma, D.P., Dhariwal, P.: Glow: generative flow with invertible 1x1 convolutions, CoRR.abs/1807.03039 (2018)
Zhu, J., Park, T., Isola, P., Efros, A.A.: Unpaired image-to-image translation using cycle-consistent adversarial networks, CoRR.abs/1703.10593 (2017)
Choi, Y., Choi, M. Kim, M., Ha, J.-W., Kim, S., Choo, J.: Stargan: Unified generative adversarial networks for multi-domain image-to-image translation, CoRR.abs/1711.09020 (2018)
Choi, Y., Uh, Y., Yoo, J., Ha, J.-W.: Stargan v2: Diverse image synthesis for multiple domains, CoRR.abs/1912.01865 (2019)
Kim, V., Kim, M., Kang, H., Lee, K.: U-gat-it: Unsupervised generative attentional networks with adaptive layer-instance normalization for image-to-image translation, CoRR.abs/1907.10830 (2019)
Ledig, C.: Photo-realistic single image super-resolution using a generative adversarial network, CoRR.abs/1609.04802 (2016)
Shaham, T.R., Dekel, T., Michaeli, T.: Singan: learning a generative model from a single natural image. In: Proceedings of ICCV (2019)
Bell-Kligler, S., Shocher, A., Irani, M.: Blind super-resolution kernel estimation using an internal-gan. In: Proceedings of NeurIPS (2019)
Clark, A., Simonyan, K., Donahue, J.: Adversarial video generation on complex datasets, CoRR.abs/1907.06571 (2019)
Arjovsky, M., Chintala, S., Bottou, L.: Wasserstein GAN, CoRR.abs/1701.07875 (2017)
Gulrajani, I., Ahmed, F., Arjovsky, M., Dumoulin, V., Courville, A.C.: Improved training of Wasserstein GANs, CoRR.abs/1704.00028 (2017)
Heusel, M., Ramsauer, H., Unterthiner, T., Nessler, B., Hochreiter, S.: GANs trained by a two time-scale update rule converge to a local Nash equilibrium. In: NIPS, pp. 6626–6637 (2017)
Miyato, T., Kataoka, T., Koyama, M., Yoshida, Y.: Spectral normalization for generative adversarial networks, CoRR.abs/1802.05957 (2018)
Sergeev, A., Balso, M.D.: Horovod: fast and easy distributed deep learning in TensorFlow, arXiv preprint arXiv:1802.05799 (2018)
Szegedy, C., Vanhoucke, V., Ioffe, S., Shlens, J., Wojna, Z.: Rethinking the inception architecture for computer vision. In: Proceedings of CVPR (2016)
Mescheder, L., Geiger, A., Nowozin, S.: Which training methods for GANs do actually converge? CoRR.abs/1801.04406 (2018)
Acknowledgment
The authors would like to thank Associate Professor Jaakko Lehtinen and Tero Karras for help with the StyleGAN2 codebase, and CSC – IT Center for Science for the GPU resources on Puhti-AI via project ID 2002415.
Author information
Authors and Affiliations
Corresponding author
Editor information
Editors and Affiliations
Rights and permissions
Copyright information
© 2021 Springer Nature Switzerland AG
About this paper
Cite this paper
Loppi, N.A., Kynkäänniemi, T. (2021). Multi-node Training for StyleGAN2. In: Del Bimbo, A., et al. Pattern Recognition. ICPR International Workshops and Challenges. ICPR 2021. Lecture Notes in Computer Science(), vol 12661. Springer, Cham. https://doi.org/10.1007/978-3-030-68763-2_51
Download citation
DOI: https://doi.org/10.1007/978-3-030-68763-2_51
Published:
Publisher Name: Springer, Cham
Print ISBN: 978-3-030-68762-5
Online ISBN: 978-3-030-68763-2
eBook Packages: Computer ScienceComputer Science (R0)
