Update README.md

Download images with ratio 3/4 to match pretrained models.

README.md

@@ -7,11 +7,14 @@
 ![cuDNN 7.3.1](https://img.shields.io/badge/cudnn-10.0-b0071e.svg?style=plastic)
 ![License CC BY-NC](https://img.shields.io/badge/license-MIT-05b502.svg?style=plastic)
 
-# GANsformer: Generative Adversarial Transformers
+# GANformer: Generative Adversarial Transformers
 <p align="center">
-  <b>Drew A. Hudson* & C. Lawrence Zitnick</b></span>
+  <b><a href="https://cs.stanford.edu/~dorarad/">Drew A. Hudson</a>* & <a href="http://larryzitnick.org/">C. Lawrence Zitnick</a></b></span>
 </p>
 
+### Update: We released the new [GANformer2 paper](https://arxiv.org/abs/2111.08960)!
+
+
 *_I wish to thank [Christopher D. Manning](https://nlp.stanford.edu/~manning/) for the fruitful discussions and constructive feedback in developing the Bipartite Transformer, especially when explored within the language representation area and also in the visual context, as well as for providing the kind financial support that allowed this work to happen!_ :sunflower:
 
 <div align="center">
@@ -20,36 +23,58 @@
 </div>
 <p></p>
 
-This is an implementation of the [GANsformer](https://arxiv.org/pdf/2103.01209.pdf) model, a novel and efficient type of transformer, explored for the task of image generation. The network employs a _bipartite structure_ that enables long-range interactions across the image, while maintaining computation of linearly efficiency, that can readily scale to high-resolution synthesis. 
+This is an implementation of the [GANformer](https://arxiv.org/pdf/2103.01209.pdf) model, a novel and efficient type of transformer, explored for the task of image generation. The network employs a _bipartite structure_ that enables long-range interactions across the image, while maintaining computation of linearly efficiency, that can readily scale to high-resolution synthesis. 
 The model iteratively propagates information from a set of latent variables to the evolving visual features and vice versa, to support the refinement of each in light of the other and encourage the emergence of compositional representations of objects and scenes. 
 In contrast to the classic transformer architecture, it utilizes multiplicative integration that allows flexible region-based modulation, and can thus be seen as a generalization of the successful StyleGAN network.
 
-**Paper**: [https://arxiv.org/pdf/2103.01209](https://arxiv.org/pdf/2103.01209)  
+<img align="right" src="https://cs.stanford.edu/people/dorarad/img3.png" width="270px">
+
+**1st Paper**: [https://arxiv.org/pdf/2103.01209](https://arxiv.org/pdf/2103.01209)  
+**2nd Paper**: [https://arxiv.org/abs/2111.08960](https://arxiv.org/abs/2111.08960)  
 **Contact**: dorarad@stanford.edu  
 **Implementation**: [`network.py`](training/network.py)
 
-### Update: All code is now ready!
-
+## Pytorch version will be released soon!
 :white_check_mark: Uploading initial code and readme  
 :white_check_mark: Image sampling and visualization script  
 :white_check_mark: Code clean-up and refacotiring, adding documentation  
 :white_check_mark: Training and data-prepreation intructions  
 :white_check_mark: Pretrained networks for all datasets  
-:white_check_mark: Extra visualizations and evaluations <!--Extra visualizations/animations and evaluation-->
+:white_check_mark: Extra visualizations and evaluations  
+⬜️ Releasing Pytorch version (coming soon!)  
+⬜️ Releasing the GANformer2 model  
+⬜️ Providing models trained for longer  
+⬜️ Releasing pre-trained models for high-resolutions (up to 1024 x 1024)  
+⬜️ Adding support for conditional image generation (will be added soon!)
 
 If you experience any issues or have suggestions for improvements or extensions, feel free to contact me either thourgh the issues page or at dorarad@stanford.edu. 
 
 ## Bibtex
 ```bibtex
-@article{hudson2021gansformer,
+@article{hudson2021ganformer,
   title={Generative Adversarial Transformers},
   author={Hudson, Drew A and Zitnick, C. Lawrence},
-  journal={arXiv preprint:2103.01209},
+  journal={Proceedings of the 38th International Conference on Machine Learning, {ICML} 2021},
+  year={2021}
+}
+
+@article{hudson2021ganformer2,
+  title={Compositional Transformers for Scene Generation},
+  author={Hudson, Drew A and Zitnick, C. Lawrence},
+  journal={Advances in Neural Information Processing Systems {NeurIPS} 2021},
   year={2021}
 }
 ```
 
+## Sample Images
+Using the pre-trained models (generated after training for ***5-7x*** less steps than StyleGAN2 models! Training our models for longer will improve the image quality further):
+<div align="center">
+  <img src="https://cs.stanford.edu/people/dorarad/samples.png" width="700px">
+</div>
+
 ## Requirements
+<img align="right" src="https://cs.stanford.edu/people/dorarad/dia.png" width="190px">
+
 - Python 3.6 or 3.7 are supported.
 - We recommend TensorFlow 1.14 which was used for development, but TensorFlow 1.15 is also supported.
 - The code was tested with CUDA 10.0 toolkit and cuDNN 7.5.
@@ -57,17 +82,21 @@ If you experience any issues or have suggestions for improvements or extensions,
 - See [`requirements.txt`](requirements.txt) for the required python packages and run `pip install -r requirements.txt` to install them.
 
 ## Quickstart & Overview
-A minimal example of using a pre-trained GANsformer can be found at [`generate.py`](generate.py). When executed, the 10-lines program downloads a pre-trained modle and uses it to generate some images:
+
+A minimal example of using a pre-trained GANformer can be found at [`generate.py`](generate.py). When executed, the 10-lines program downloads a pre-trained modle and uses it to generate some images:
 ```python
-python generate.py --gpus 0 --model gdrive:bedrooms-snapshot.pkl --output-dir images --images-num 8
+python generate.py --gpus 0 --model gdrive:bedrooms-snapshot.pkl --output-dir images --images-num 32
 ```
-You can use `--truncation-psi` to control the generated images quality/diversity trade-off.
+You can use `--truncation-psi` to control the generated images quality/diversity trade-off.  
+We recommend setting it to values in the range of `0.6-1.0`.
+
+We currently provide pretrained models for resolution 256&times;256 but keep training them and will release newer checkpoints as well as pretrained models for resolution 1024&times;1024 soon!
 
 We can train and evaluate new or pretrained model both quantitatively and qualitative with [`run_netowrk.py`](run_network.py).  
 The model architecutre can be found at [`network.py`](training/network.py). The training procedure is implemented at [`training_loop.py`](training/training_loop.py).
 
 ## Data preparation
-We explored the GANsformer model on 4 datasets for images and scenes: [CLEVR](https://cs.stanford.edu/people/jcjohns/clevr/), [LSUN-Bedrooms](https://www.yf.io/p/lsun), [Cityscapes](https://www.cityscapes-dataset.com/) and [FFHQ](https://github.com/NVlabs/ffhq-dataset). The model can be trained on other datasets as well.
+We explored the GANformer model on 4 datasets for images and scenes: [CLEVR](https://cs.stanford.edu/people/jcjohns/clevr/), [LSUN-Bedrooms](https://www.yf.io/p/lsun), [Cityscapes](https://www.cityscapes-dataset.com/) and [FFHQ](https://github.com/NVlabs/ffhq-dataset). The model can be trained on other datasets as well.
 We trained the model on `256x256` resolution. Higher resolutions are supported too. The model will automatically adapt to the resolution of the images in the dataset.
 
 The [`prepare_data.py`](prepare_data.py) can either prepare the datasets from our catalog or create new datasets.
@@ -75,7 +104,7 @@ The [`prepare_data.py`](prepare_data.py) can either prepare the datasets from ou
 ### Default Datasets 
 To prepare the datasets from the catalog, run the following command:
 ```python
-python prepare_data.py --ffhq --cityscapes --clevr --bedroom --max-images 100000
+python prepare_data.py --ffhq --cityscapes --clevr --bedrooms --max-images 100000
 ```
 
 See table below for details about the datasets in the catalog.
@@ -104,15 +133,16 @@ The script supports several formats: `png`, `jpg`, `npy`, `hdf5`, `tfds` and `lm
 Use `--max-images` to reduce the size of the `tfrecord` files.
 
 ## Training
-Models are trained by using the `--train` option. To fine-tune a pretrained GANsformer model:
+Models are trained by using the `--train` option. To fine-tune a pretrained GANformer model:
 ```python
-python run_network.py --train --gpus 0 --gansformer-default --expname clevr-pretrained --dataset clevr  
+python run_network.py --train --gpus 0 --ganformer-default --expname clevr-pretrained --dataset clevr \
+  --pretrained-pkl gdrive:clevr-snapshot.pkl
 ```
+We provide pretrained models for `bedrooms`, `cityscapes`, `clevr` and `ffhq`.
 
-To train a GANsformer in its default configuration form scratch:
+To train a GANformer in its default configuration form scratch:
 ```python
-python run_network.py --train --gpus 0 --gansformer-default --expname clevr-scratch --dataset clevr \
-  --pretrained-pkl None
+python run_network.py --train --gpus 0 --ganformer-default --expname clevr-scratch --dataset clevr
 ```
 
 By defualt, models training is resumed from the latest snapshot. Use `--restart` to strat a new experiment, or `--pretrained-pkl` to select a particular snapshot to load.
@@ -121,13 +151,15 @@ For comparing to state-of-the-art, we compute metric scores using 50,000 sample
 
 We support a large variety of command-line options to adjust the model, training, and evaluation. Run `python run_network.py -h` for the full list of options!
 
+we recommend exploring different values for `--gamma` when training on new datasets. If you train on resolution >= 512 and observe OOM issues, consider reducing `--minibatch-size` to a lower value.
+
 ### Logging
 * During training, sample images and attention maps will be generated and stored at results/<expname>-<run-id> (`--keep-samples`).
 * Metrics will also be regularly commputed and reported in a `metric-<name>.txt` file. `--metrics` can be set to `fid` for FID, `is` for Inception Score and `pr` for Precision/Recall.
 * Tensorboard logs are also created (`--summarize`) that track the metrics, loss values for the generator and discriminator, and other useful statistics over the course of training.
 
 ### Baseline models
-The codebase suppors multiple baselines in addition to the GANsformer. For instance, to run a vanilla GAN model:
+The codebase suppors multiple baselines in addition to the GANformer. For instance, to run a vanilla GAN model:
 ```python
 python run_network.py --train --gpus 0 --baseline GAN --expname clevr-gan --dataset clevr 
 ```
@@ -141,9 +173,9 @@ To evalute a model, use the `--eval` option:
 ```python
 python run_network.py --eval --gpus 0 --expname clevr-exp --dataset clevr
 ```
-Add `--pretrained-network gdrive:<dataset>-snapshot.pkl` to evalute a pretrained model.
+Add `--pretrained-pkl gdrive:<dataset>-snapshot.pkl` to evalute a pretrained model.
 
-Below we provide the FID-50k scores for the GANsformer (_using the pretrained checkpoints above_) as well as baseline models.  
+Below we provide the FID-50k scores for the GANformer (_using the pretrained checkpoints above_) as well as baseline models.  
 Note that these scores are different than the scores reported in the StyleGAN2 paper since they run experiments for up to 7x more training steps (5k-15k kimg-steps in our experiments over all models, which takes about 3-4 days with 4 GPUs, vs 50-70k kimg-steps in their experiments, which take over 90 GPU-days).
 
 | Model          | CLEVR        | LSUN-Bedroom | FFHQ       | Cityscapes |
@@ -153,20 +185,27 @@ Note that these scores are different than the scores reported in the StyleGAN2 p
 | **SAGAN**      | 26.04        | 14.06        | 16.21      | 12.81      |
 | **StyleGAN2**  | 16.05        | 11.53        | 16.21      | 8.35       |
 | **VQGAN**      | 32.60        | 59.63        | 63.12      | 173.80     |
-| **GANsformer** | ***9.24***   | ***6.15***   | ***7.42*** | ***5.23*** |
+| **GANformer** | ***9.24***   | ***6.15***   | ***7.42*** | ***5.23*** |
 
 <div>
   <img src="https://cs.stanford.edu/people/dorarad/plot1.png" width="350px">
   <img src="https://cs.stanford.edu/people/dorarad/plot2.png" width="350px">
 </div>
 
+### Model Change-log
+Compared to the original GANformer depicted in the paper, this repository make several additional improvments that contributed to the performance:
+* Use `--mapping_ltnt2ltnt` so that the latents communicate with each other directly through self-attention inside the mapping network before starting to generate the image.
+* Add an additional global latent (`--style`) to the `k` latent components, such that first the global latent modulates all the image features uniformly, and then the `k` latents modulate different regions based on the bipartite transformer's attention.  
+The global latent is useful for coordinating holistic aspects of the image such as global lighting conditions, global style properties for e.g. faces, etc.
+* After making these changes, we observed no additional benefit from adding the transformer to the discriminator, and therefore for simplicity we disabled that.
+
 ## Visualization
 The code supports producing qualitative results and visualizations. For instance, to create attention maps for each layer:
 ```python
 python run_network.py --gpus 0 --eval --expname clevr-exp --dataset clevr --vis-layer-maps
 ```
 
-Below you can see sample images and attention maps produced by the GANsformer:
+Below you can see sample images and attention maps produced by the GANformer:
 
 <div align="center">
   <img src="https://cs.stanford.edu/people/dorarad/atts.png" style="float:left" width="831px">
@@ -176,15 +215,15 @@ Below you can see sample images and attention maps produced by the GANsformer:
 In the following we list some of the most useful model options. 
 
 ### Training
-* `--gamma`: We recommend explore different values for the chosen dataset (default: `10`)
-* `--truncation-psi`: Controls the image quality/diversity trade-off. (default: `0.65`)
+* `--gamma`: We recommend exploring different values for the chosen dataset (default: `10`)
+* `--truncation-psi`: Controls the image quality/diversity trade-off. (default: `0.7`)
 * `--eval-images-num`: Number of images to compute metrics over. We recommend selecting a lower number to expedite training (default: `50,000`)
 * `--restart`: To restart training from sracth instead of resuming from the latest snapshot
 * `--pretrained-pkl`: To load a pretrained model, either a local one or from drive `gdrive:<dataset>-snapshot.pkl` for the datasets in the catalog.
 * `--data-dir` and `--result-dir`: Directory names for the datasets (`tfrecords`) and logging/results.
 
 ### Model (most useful)
-* `--transformer`: To add transformer layers to the generator (GANsformer)
+* `--transformer`: To add transformer layers to the generator (GANformer)
 * `--components-num`: Number of latent components, which will attend to the image. We recommend values in the range of `8-16` (default: `1`)
 * `--latent-size`: Overall latent size (default: `512`). The size of each latent component will then be `latent_size/components_num`
 * `--num-heads`: Number of attention heads (default: `1`)
@@ -214,16 +253,21 @@ In the following we list some of the most useful model options.
 
 Run `python run_network.py -h` for the full options list.
 
-## Sample images
+## Sample images (more examples)
 <div align="center">
-  <img src="https://cs.stanford.edu/people/dorarad/bedroom.png" style="float:left" width="420px">
-  <img src="https://cs.stanford.edu/people/dorarad/clevr.png" style="float:left" width="411px">
+  <img src="https://cs.stanford.edu/people/dorarad/faces.png" style="float:left" width="750px">
+  <br>
+  <img src="https://cs.stanford.edu/people/dorarad/bedroom.png" style="float:left" width="750px">
+  <br>
+  <img src="https://cs.stanford.edu/people/dorarad/clevr_new.png" style="float:left" width="750px">
+  <br>
+  <img src="https://cs.stanford.edu/people/dorarad/cities_small.png" style="float:left" width="750px">
 </div>
 
 ## CUDA / Installation
 The model relies on custom TensorFlow ops that are compiled on the fly using [NVCC](https://docs.nvidia.com/cuda/cuda-compiler-driver-nvcc/index.html). 
 
-To set up the environment:
+To set up the environment e.g. for cuda-10.0:
 ```python
 export PATH=/usr/local/cuda-10.0/bin${PATH:+:${PATH}}
 export LD_LIBRARY_PATH=/usr/local/cuda10.0/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
@@ -237,11 +281,11 @@ nvcc test_nvcc.cu -o test_nvcc -run
 ```
 
 ## Architecture Overview
-The GANsformer consists of two networks:
+The GANformer consists of two networks:
 
-**Generator**: which produces the images (`x`) given randomly sampled latents (`z`). The latent z has a shape `[batch_size, component_num, latent_dim]`, where `component_num = 1` by default (Vanilla GAN, StyleGAN) but is > 1 for the GANsformer model. We can define the latent components by splitting `z` along the second dimension to obtain `z_1,...,z_k` latent components. The generator likewise consists of two parts:
+**Generator**: which produces the images (`x`) given randomly sampled latents (`z`). The latent z has a shape `[batch_size, component_num, latent_dim]`, where `component_num = 1` by default (Vanilla GAN, StyleGAN) but is > 1 for the GANformer model. We can define the latent components by splitting `z` along the second dimension to obtain `z_1,...,z_k` latent components. The generator likewise consists of two parts:
 * **Mapping network**: converts sampled latents from a normal distribution (`z`) to the intermediate space (`w`). A series of Feed-forward layers. The k latent components either are mapped independently from the `z` space to the `w` space or interact with each other through self-attention (optional flag).
-* **Synthesis network**: the intermediate latents w are used to guide the generation of new images. Images features begin from a small constant/sampled grid of `4x4`, and then go through multiple layers of convolution and up-sampling until reaching the desirable resolution (e.g. `256x256`). After each convolution, the image features are modulated (meaning that their variance and bias are controlled) by the intermediate latent vectors `w`. While in the StyleGAN model there is one global w vectors that controls all the features equally. The GANsformer uses attention so that the k latent components specialize to control different regions in the image to create it cooperatively, and therefore perform better especially in generating images depicting multi-object scenes.
+* **Synthesis network**: the intermediate latents w are used to guide the generation of new images. Images features begin from a small constant/sampled grid of `4x4`, and then go through multiple layers of convolution and up-sampling until reaching the desirable resolution (e.g. `256x256`). After each convolution, the image features are modulated (meaning that their variance and bias are controlled) by the intermediate latent vectors `w`. While in the StyleGAN model there is one global w vectors that controls all the features equally. The GANformer uses attention so that the k latent components specialize to control different regions in the image to create it cooperatively, and therefore perform better especially in generating images depicting multi-object scenes.
 * **Attention** can be used in several ways
   * **Simplex Attention**: when attention is applied in one direction only from the latents to the image features (**top-down**).
   * **Duplex Attention**: when attention is applied in the two directions: latents to image features (**top-down**) and then image features back to latents (**bottom-up**), so that each representation informs the other iteratively.
@@ -253,6 +297,6 @@ The GANsformer consists of two networks:
 ## Codebase
 This codebase builds on top of and extends the great [StyleGAN2 repository](https://github.com/NVlabs/stylegan2) by Karras et al.  
 
-The GANsformer model can also be seen as a generalization of StyleGAN: while StyleGAN has one global latent vector that control the style of all image features globally, the GANsformer has *k* latent vectors, that cooperate through attention to control regions within the image, and thereby better modeling images of multi-object and compositional scenes.
+The GANformer model can also be seen as a generalization of StyleGAN: while StyleGAN has one global latent vector that control the style of all image features globally, the GANformer has *k* latent vectors, that cooperate through attention to control regions within the image, and thereby better modeling images of multi-object and compositional scenes.
 
 If you have questions, comments or feedback, please feel free to contact me at dorarad@stanford.edu, Thank you! :)

generate.py

@@ -1,4 +1,4 @@
-# Minimal script for generating images using pre-trained the GANsformer
+# Minimal script for generating images using pre-trained the GANformer
 # Ignore all future warnings
 from warnings import simplefilter
 simplefilter(action = "ignore", category = FutureWarning)
@@ -9,11 +9,13 @@ import numpy as np
 from tqdm import tqdm
 
 from training import misc
+from training.misc import crop_max_rectangle as crop
+
 import dnnlib.tflib as tflib
 from pretrained_networks import load_networks # returns G, D, Gs
 # G: generator, D: discriminator, Gs: generator moving-average (higher quality images)
 
-def run(model, gpus, output_dir, images_num, truncation_psi, batch_size):
+def run(model, gpus, output_dir, images_num, truncation_psi, batch_size, ratio):
     print("Loading networks...")
     os.environ["CUDA_VISIBLE_DEVICES"] = gpus                   # Set GPUs
     tflib.init_tf()                                             # Initialize TensorFlow
@@ -29,18 +31,19 @@ def run(model, gpus, output_dir, images_num, truncation_psi, batch_size):
     os.makedirs(output_dir, exist_ok = True)                    # Make output directory
     pattern = "{}/Sample_{{:06d}}.png".format(output_dir)       # Output images pattern
     for i, image in tqdm(list(enumerate(images))):              # Save images
-        misc.to_pil(image).save(pattern.format(i))
+        crop(misc.to_pil(image), ratio).save(pattern.format(i))
 
 def main():
-    parser = argparse.ArgumentParser(description = "Generate images with the GANsformer")
+    parser = argparse.ArgumentParser(description = "Generate images with the GANformer")
     parser.add_argument("--model",              help = "Filename for a snapshot to resume (optional)", default = None, type = str)
     parser.add_argument("--gpus",               help = "Comma-separated list of GPUs to be used (default: %(default)s)", default = "0", type = str)
     parser.add_argument("--output-dir",         help = "Root directory for experiments (default: %(default)s)", default = "images", metavar = "DIR")
-    parser.add_argument("--images-num",         help = "Number of images to generate (default: %(default)s)", default = 10, type = int)
-    parser.add_argument("--truncation-psi",     help = "Truncation Psi to be used in producing sample images (default: %(default)s)", default = 0.5, type = float)
+    parser.add_argument("--images-num",         help = "Number of images to generate (default: %(default)s)", default = 32, type = int)
+    parser.add_argument("--truncation-psi",     help = "Truncation Psi to be used in producing sample images (default: %(default)s)", default = 0.7, type = float)
     parser.add_argument("--batch-size",         help = "Batch size for generating images (default: %(default)s)", default = 8, type = int)
+    parser.add_argument("--ratio",              help = "Crop ratio for output images (default: %(default)s)", default = 1.0, type = float)
     args = parser.parse_args()
     run(**vars(args))
 
 if __name__ == "__main__":
-    main()
+    main()

metrics/metric_base.py

@@ -11,6 +11,7 @@ import dnnlib.tflib as tflib
 
 from training import misc
 from training import dataset
+from tqdm import tqdm
 
 # Base class for metrics
 class MetricBase:
@@ -68,10 +69,11 @@ class MetricBase:
 
     def run(self, network_pkl, num_imgs, run_dir = None, data_dir = None, dataset_args = None,
             mirror_augment = None, num_gpus = 1, tf_config = None, log_results = True,
-            Gs_kwargs = dict(is_validation = True), **kwargs):
+            Gs_kwargs = dict(is_validation = True), eval_mod = False, **kwargs):
 
         self._reset(network_pkl = network_pkl, run_dir = run_dir, data_dir = data_dir,
             dataset_args = dataset_args, mirror_augment = mirror_augment)
+        self.eval_mod = eval_mod
 
         time_begin = time.time()
         with tf.Graph().as_default(), tflib.create_session(tf_config).as_default():
@@ -189,8 +191,10 @@ class MetricBase:
     # Returns the features [num_imgs, featurizer.output_shape[1]]
     def _get_feats(self, img_iter, featurizer, minibatch_size, num_gpus, num_imgs):
         feats = np.empty([num_imgs, featurizer.output_shape[1]], dtype = np.float32)
-
-        for idx, imgs in enumerate(img_iter):
+        itr = enumerate(img_iter)
+        if self.eval_mod:
+            itr = tqdm(itr, total = num_imgs / minibatch_size, unit_scale = minibatch_size)
+        for idx, imgs in itr:
             begin = idx * minibatch_size
             end = min(begin + minibatch_size, num_imgs)
 
@@ -225,8 +229,12 @@ class MetricBase:
                 result_expr.append(featurizer.clone().get_output_for(imgs))
 
         # Compute features for newly generated 'num_imgs' images
+        itr = range(0, num_imgs, minibatch_size)
+        if self.eval_mod:
+            itr = tqdm(itr, total = num_imgs / minibatch_size, unit_scale = minibatch_size)
+
         feats = np.empty([num_imgs, featurizer.output_shape[1]], dtype = np.float32)
-        for begin in range(0, num_imgs, minibatch_size):
+        for begin in itr:
             self._report_progress(begin, num_imgs)
             end = min(begin + minibatch_size, num_imgs)
             feats[begin:end] = np.concatenate(tflib.run(result_expr), axis = 0)[:end-begin]

metrics/precision_recall.py

@@ -1,4 +1,4 @@
-# Precision/Recall (PR).
+# Precision/Recall (PR)
 import os
 import numpy as np
 import tensorflow as tf
@@ -120,7 +120,7 @@ class ManifoldEstimator():
         return batch_predictions
 
 def knn_precision_recall_features(ref_features, eval_features, feature_net, nhood_sizes,
-                                  row_batch_size, col_batch_size, num_gpus):
+                                  row_batch_size, col_batch_size, num_imgs, num_gpus):
     # Computes k-NN precision and recall for two sets of feature vectors
     state = dnnlib.EasyDict()
     num_features = feature_net.output_shape[1]
@@ -173,14 +173,14 @@ class PR(metric_base.MetricBase):
 
         if paths is not None:
             # Extract features for local sample image files (paths)
-            feats = self._paths_to_feats(paths, feat_func, minibatch_size, num_gpus, num_imgs)
+            eval_features = self._paths_to_feats(paths, feature_net, minibatch_size, num_gpus, num_imgs)
         else:
             # Extract features for newly generated fake imgs
-            feats = self._gen_feats(Gs, inception, minibatch_size, num_gpus, Gs_kwargs)
+            eval_features = self._gen_feats(Gs, feature_net, minibatch_size, num_imgs, num_gpus, Gs_kwargs)
 
         # Compute precision and recall
         state = knn_precision_recall_features(ref_features = ref_features, eval_features = eval_features,
             feature_net = feature_net, nhood_sizes = [self.nhood_size], row_batch_size = self.row_batch_size,
-            col_batch_size = self.row_batch_size, num_gpus = num_gpus)
+            col_batch_size = self.row_batch_size, num_imgs = num_imgs, num_gpus = num_gpus)
         self._report_result(state.knn_precision[0], suffix = "_precision")
         self._report_result(state.knn_recall[0], suffix = "_recall")

prepare_data.py

@@ -24,25 +24,13 @@ catalog = {
    "ffhq": EasyDict({
         "name": "FFHQ", # Dataset name for logging 
         "filename": "ffhq-r08.tfrecords1of1", # Local file name
-        "url": "https://drive.google.com/uc?id=1Oo-43djEakn6AzeZidfmHWYFNfMen8ZP", # download URL
+        "url": "http://downloads.cs.stanford.edu/nlp/data/dorarad/ffhq-r08.tfrecords1of1", # download URL
         "md5": "74de4f07dc7bfb07c0ad4471fdac5e67", # MD5 checksum to potentially skip download
         "ratio": 1, # height/width ratio
         "size": 13, # download size in GB
         "shards": 1, # Number of tfrecord shards
         "img_num": 70000 # Number of images
     }),
-    "clevr": EasyDict({
-        "name": "CLEVR", # Dataset name for logging 
-        "filename": "CLEVR_v1.0.zip",
-        "url": "https://dl.fbaipublicfiles.com/clevr/CLEVR_v1.0.zip",
-        "dir": "CLEVR_v1.0/images", # Image directory to process while turning into tfrecords
-        "md5": "b11922020e72d0cd9154779b2d3d07d2",
-        "ratio": 0.75,
-        "size": 18,
-        "shards": 5,
-        "img_num": 100000,
-        "process": dataset_tool.create_from_imgs # Function to convert download to tfrecords
-    }),
     "bedrooms": EasyDict({
         "name": "LSUN-Bedrooms", # Dataset name for logging 
         "filename": "bedroom_train_lmdb.zip",
@@ -64,7 +52,26 @@ catalog = {
         "size": 2,
         "shards": 16,
         "img_num": 25000
-    })
+    }),
+    "clevr": EasyDict({
+        "name": "CLEVR", # Dataset name for logging 
+        "filename": "clevr.zip",
+        "url": "https://drive.google.com/uc?id=1lY4JE30yk26v0MWHNpXBOMzltufUcTXj",
+        "md5": "3040bb20a29cd2f0e1e9231aebddf2a1",
+        "size": 6,
+        "ratio": 0.75,
+        "shards": 5,
+        "img_num": 100000
+        ##########################################################################################
+        # Currently, we download preprocessed TFrecords of CLEVR images with image ratio 0.75.
+        # To process instead the dataset from scratch (with the original image ratio of 320/480), add the following:
+        # "filename": "CLEVR_v1.0.zip",
+        # "size": 18,
+        # "dir": "CLEVR_v1.0/images", # Image directory to process while turning into tfrecords
+        # "url": "https://dl.fbaipublicfiles.com/clevr/CLEVR_v1.0.zip",
+        # "md5": "b11922020e72d0cd9154779b2d3d07d2",
+        # "process": dataset_tool.create_from_imgs # Function to convert download to tfrecords
+    })    
 }
 
 formats_catalog = {
@@ -141,22 +148,22 @@ def prepare(tasks, data_dir, shards_num = 1, max_images = None,
     mkdir(data_dir)
     for task in tasks:
         # If task not in catalog, create custom task configuration
-        c = catalog.get(task, {
+        c = catalog.get(task, EasyDict({
             "local": True,
             "name": task,
             "dir": images_dir,
             "ratio": ratio,
             "process": formats_catalog.get(format)
-        })
+        }))
 
         dirname = "{}/{}".format(data_dir, task)
         mkdir(dirname)
 
         # try:
         print(misc.bold("Preparing the {} dataset...".format(c.name)))
-        fname = "{}/{}".format(dirname, c.filename)
-
+        
         if "local" not in c:
+            fname = "{}/{}".format(dirname, c.filename)
             download = not ((os.path.exists(fname) and verify_md5(fname, c.md5)))
             path = get_path(c.url, dirname, path = c.filename)
 
@@ -183,15 +190,15 @@ def prepare(tasks, data_dir, shards_num = 1, max_images = None,
         #     print(sys.exc_info())
 
 def run_cmdline(argv):
-    parser = argparse.ArgumentParser(prog = argv[0], description = "Download and prepare data for the GANsformer.")
+    parser = argparse.ArgumentParser(prog = argv[0], description = "Download and prepare data for the GANformer.")
     parser.add_argument("--data-dir",       help = "Directory of created dataset", default = "datasets", type = str)
     parser.add_argument("--shards-num",     help = "Number of shards to split each dataset to (optional)", default = 1, type = int)
     parser.add_argument("--max-images",     help = "Maximum number of images to have in the dataset (optional). Use to reduce the produced tfrecords file size", default = None, type = int)
     # Default tasks
-    parser.add_argument("--clevr",          help = "Prepare the CLEVR dataset (18GB download, up to 15.5GB tfrecords, 100k images)", dest = "tasks", action = "append_const", const = "clevr")
-    parser.add_argument("--bedrooms",       help = "Prepare the LSUN-bedrooms dataset (42.8GB, up to 480GB tfrecords, 3M images)", dest = "tasks", action = "append_const", const = "bedrooms")
+    parser.add_argument("--clevr",          help = "Prepare the CLEVR dataset (6.41GB download, 31GB tfrecords, 100k images)", dest = "tasks", action = "append_const", const = "clevr")
+    parser.add_argument("--bedrooms",       help = "Prepare the LSUN-bedrooms dataset (42.8GB download, up to 480GB tfrecords, 3M images)", dest = "tasks", action = "append_const", const = "bedrooms")
     parser.add_argument("--ffhq",           help = "Prepare the FFHQ dataset (13GB download, 13GB tfrecords, 70k images)", dest = "tasks", action = "append_const", const = "ffhq")
-    parser.add_argument("--cityscapes",     help = "Prepare the cityscapes dataset (1.8GB, 8GB tfrecords, 25k images)", dest = "tasks", action = "append_const", const = "cityscapes")
+    parser.add_argument("--cityscapes",     help = "Prepare the cityscapes dataset (1.8GB download, 8GB tfrecords, 25k images)", dest = "tasks", action = "append_const", const = "cityscapes")
     # Create a new task with custom images
     parser.add_argument("--task",           help = "New dataset name", type = str, dest = "tasks", action = "append")
     parser.add_argument("--images-dir",     help = "Provide source image directory to convert into tfrecords (will be searched recursively)", default = None, type = str)
@@ -200,10 +207,11 @@ def run_cmdline(argv):
 
     args = parser.parse_args()
     if not args.tasks:
-        print("No tasks specified. Please see '-h' for help.")
-        exit(1)
+        misc.error("No tasks specified. Please see '-h' for help.")
+    if args.max_images < 50000:
+        misc.log("Warning: max-images is set to {}. We recommend setting it at least to 50,000 to allow statistically correct computation of the FID-50k metric.".format(args.max_images), "red")
 
     prepare(**vars(args))
 
 if __name__ == "__main__":
-    run_cmdline(sys.argv)
+    run_cmdline(sys.argv)

pretrained_networks.py

@@ -6,7 +6,7 @@ import dnnlib.tflib as tflib
 gdrive_urls = {
     "gdrive:clevr-snapshot.pkl": "https://drive.google.com/uc?id=1JNb9r2j0MKpygxn-FT-HGClfzZ0CV_dy",
     "gdrive:cityscapes-snapshot.pkl": "https://drive.google.com/uc?id=1SLImpJEw--fKe9ngM99-KKXOXmmoWw6g",
-    "gdrive:ffhq-snapshot.pkl": "https://drive.google.com/uc?id=1JSZLVmhxOhUsC4Ijh6PmVoSIVsJTwyR7",
+    "gdrive:ffhq-snapshot.pkl": "https://drive.google.com/uc?id=13rF5RXN4-FeEZX9Ph40jXnt7Bqw0PJDk",
     "gdrive:bedrooms-snapshot.pkl": "https://drive.google.com/uc?id=1-2L3iCBpP_cf6T2onf3zEQJFAAzxsQne"
 }
 

run_network.py

@@ -27,7 +27,8 @@ def cset(dicts, name, prop):
 
 # Conditional set: if dict[name] is not populated from the command line, then assign dict[name] := prop
 def nset(args, name, prop):
-    if name not in sys.argv:
+    flag = "--{}".format(name.replace("_", "-"))
+    if flag not in sys.argv:
         args[name] = prop
 
 # Conditional set: if dict[name] has its default value, then assign dict[name] := prop
@@ -38,7 +39,7 @@ def dset(d, name, prop, default):
 # Set network (generator or discriminator): model, loss and optimizer
 def set_net(net, reg_interval):
     ret = EasyDict()
-    ret.args  = EasyDict(func_name = "training.network.{}_GANsformer".format(net[0])) # network options
+    ret.args  = EasyDict(func_name = "training.network.{}_GANformer".format(net[0])) # network options
     ret.loss_args = EasyDict(func_name = "training.loss.{}_loss".format(net[0]))      # loss options
     ret.opt_args  = EasyDict(beta1 = 0.0, beta2 = 0.99, epsilon = 1e-8)               # optimizer options
     ret.reg_interval = reg_interval
@@ -60,14 +61,10 @@ def run(**args):
     if not args.train and not args.eval:
         misc.log("Warning: Neither --train nor --eval are provided. Therefore, we only print network shapes", "red")
 
-    if args.gansformer_default:
+    if args.ganformer_default:
         task = args.dataset
-        pretrained = "gdrive:{}-snapshot.pkl".format(task)
-        if pretrained not in pretrained_networks.gdrive_urls:
-            pretrained = None
 
-        nset(args, "recompile", pretrained is not None)
-        nset(args, "pretrained_pkl", pretrained)
+        nset(args, "recompile", args.pretrained_pkl is not None)
         nset(args, "mirror_augment", task in ["cityscapes", "ffhq"])
 
         nset(args, "transformer", True)
@@ -86,7 +83,7 @@ def run(**args):
 
         gammas = {
             "ffhq": 10, 
-            "cities": 20, 
+            "cityscapes": 20, 
             "clevr": 40, 
             "bedrooms": 100
         }
@@ -134,7 +131,7 @@ def run(**args):
     }
     args.ratio = ratios.get(args.dataset, args.ratio)
     dataset_args = EasyDict(tfrecord_dir = args.dataset, max_imgs = args.train_images_num, 
-        num_threads = args.num_threads)
+        num_threads = args.num_threads, resolution = args.resolution)
     for arg in ["data_dir", "mirror_augment", "total_kimg", "ratio"]:
         cset(train, arg, args[arg])
 
@@ -197,7 +194,7 @@ def run(**args):
     if args.components_num > 1:
         if not (args.transformer or args.kgan):
             misc.error("--components-num > 1 but the model is not using components. " + 
-                "Either add --transformer for GANsformer or --kgan for k-GAN.")
+                "Either add --transformer for GANformer or --kgan for k-GAN.")
 
         args.latent_size = int(args.latent_size / args.components_num)
     cD.args.latent_size = cG.args.latent_size = cG.args.dlatent_size = args.latent_size
@@ -213,7 +210,7 @@ def run(**args):
         cset(cG.args, arg, args[arg])
     cD.args.mbstd_group_size = args.minibatch_std_size
 
-    # GANsformer
+    # GANformer
     cset(cG.args, "transformer", args.transformer)
     cset(cD.args, "transformer", args.d_transformer)
 
@@ -237,7 +234,7 @@ def run(**args):
     cset([cG.args, train], "merge", args.kgan)
 
     if args.kgan and args.transformer:
-        misc.error("Either have --transformer for GANsformer or --kgan for k-GAN, not both")
+        misc.error("Either have --transformer for GANformer or --kgan for k-GAN, not both")
 
     # Attention
     for arg in ["start_res", "end_res", "ltnt2ltnt", "img2img"]: # , "local_attention"
@@ -354,7 +351,7 @@ def _parse_comma_sep(s):
     return s.split(",")
 
 def main():
-    parser = argparse.ArgumentParser(description = "Train the GANsformer")
+    parser = argparse.ArgumentParser(description = "Train the GANformer")
 
     # Framework
     # ------------------------------------------------------------------------------------------------------
@@ -364,7 +361,7 @@ def main():
     parser.add_argument("--gpus",               help = "Comma-separated list of GPUs to be used (default: %(default)s)", default = "0", type = str)
 
     ## Default configurations
-    parser.add_argument("--gansformer-default", help = "Select a default GANsformer configuration, either pretrained (default) or from scratch (with --pretrained-pkl None)", default = None, action = "store_true")
+    parser.add_argument("--ganformer-default", help = "Select a default GANformer configuration, either pretrained (default) or from scratch (with --pretrained-pkl None)", default = None, action = "store_true")
     parser.add_argument("--baseline",           help = "Use a baseline model configuration", default = None, choices = ["GAN", "StyleGAN2", "kGAN", "SAGAN"], type = str)
 
     ## Resumption
@@ -380,6 +377,7 @@ def main():
     parser.add_argument("--data-dir",           help = "Datasets root directory (default: %(default)s)", default = "datasets", metavar = "DIR")
     parser.add_argument("--dataset",            help = "Training dataset name (subdirectory of data-dir)", required = True)
     parser.add_argument("--ratio",              help = "Image height/width ratio in the dataset", default = 1.0, type = float)
+    parser.add_argument("--resolution",         help = "Training resolution", default = 256, type = int)
     parser.add_argument("--num-threads",        help = "Number of input processing threads (default: %(default)s)", default = 4, type = int)
     parser.add_argument("--mirror-augment",     help = "Perform horizontal flip augmentation for the data (default: %(default)s)", default = None, action = "store_true")
     parser.add_argument("--train-images-num",   help = "Maximum number of images to train on. If not specified, train on the whole dataset.", default = None, type = int)
@@ -405,8 +403,8 @@ def main():
     ## Visualization
     parser.add_argument("--vis-images",         help = "Save image samples", default = None, action = "store_true")
     parser.add_argument("--vis-latents",        help = "Save latent vectors", default = None, action = "store_true")
-    parser.add_argument("--vis-maps",           help = "Save attention maps (for GANsformer only)", default = None, action = "store_true")
-    parser.add_argument("--vis-layer-maps",     help = "Save attention maps for all layers (for GANsformer only)", default = None, action = "store_true")
+    parser.add_argument("--vis-maps",           help = "Save attention maps (for GANformer only)", default = None, action = "store_true")
+    parser.add_argument("--vis-layer-maps",     help = "Save attention maps for all layers (for GANformer only)", default = None, action = "store_true")
     parser.add_argument("--vis-interpolations", help = "Create latent interpolations", default = None, action = "store_true")
     parser.add_argument("--vis-noise-var",      help = "Create noise variation visualization", default = None, action = "store_true")
     parser.add_argument("--vis-style-mix",      help = "Create style mixing visualization", default = None, action = "store_true")
@@ -457,7 +455,7 @@ def main():
     parser.add_argument("--local-noise",        help = "Add stochastic local noise each layer (default: %(default)s)", default = True, metavar = "BOOL", type = _str_to_bool, nargs = "?")
     parser.add_argument("--minibatch-std-size", help = "Add minibatch standard deviation layer in the discriminator, 0 to disable (default: %(default)s)", default = 4, type = int)
 
-    ## GANsformer
+    ## GANformer
     parser.add_argument("--transformer",        help = "Add transformer layers to the generator: top-down latents-to-image (default: False)", default = None, action = "store_true")
     parser.add_argument("--latent-size",        help = "Latent size, summing the dimension of all components (default: %(default)s)", default = 512, type = int)
     parser.add_argument("--components-num",     help = "Components number. Each component has latent dimension of 'latent-size / components-num'. " +

training/dataset.py

@@ -76,13 +76,18 @@ class TFRecordDataset:
         # Determine shape and resolution
         max_shape = max(tfr_shapes, key = np.prod)
         self.resolution = resolution if resolution is not None else max_shape[1]
+        if resolution is not None:
+            assert resolution <= max_shape[1]
         self.resolution_log2 = int(np.log2(self.resolution))
+        file_indexes = [i for i, tfr_shape in enumerate(tfr_shapes) if tfr_shape[1] == self.resolution]
+        tfr_files = [tfr_files[i] for i in file_indexes]
+        tfr_shapes = [tfr_shapes[i] for i in file_indexes]
         self.shape = [max_shape[0], self.resolution, self.resolution]
         tfr_lods = [self.resolution_log2 - int(np.log2(shape[1])) for shape in tfr_shapes]
         assert all(shape[0] == max_shape[0] for shape in tfr_shapes)
         assert all(shape[1] == shape[2] for shape in tfr_shapes)
         assert all(shape[1] == self.resolution // (2**lod) for shape, lod in zip(tfr_shapes, tfr_lods))
-        assert all(lod in range(self.resolution_log2 - 1) for lod in tfr_lods)
+        #assert all(lod in range(self.resolution_log2 - 1) for lod in tfr_lods)
 
         # Load labels
         assert max_label_size == "full" or max_label_size >= 0

training/misc.py

@@ -46,7 +46,7 @@ def log(txt, color = None):
 
 def open_file_or_url(file_or_url):
     if dnnlib.util.is_url(file_or_url):
-        return dnnlib.util.open_url(file_or_url, cache_dir = ".GANsformer-cache")
+        return dnnlib.util.open_url(file_or_url, cache_dir = ".GANformer-cache")
     return open(file_or_url, "rb")
 
 def load_pkl(file_or_url):
@@ -134,7 +134,7 @@ def crop_center(img, cw, ch):
 def crop_max_rectangle(img, ratio = 1.0):
     if ratio is None:
         return img
-    s = min(img.size)
+    s = min(img.size[0], img.size[1] / ratio)
     return crop_center(img, s, ratio * s)
 
 # Pad an image of dimension w,h to the smallest containing square (max(w,h), max(w,h))
@@ -358,7 +358,7 @@ def save_gif(imgs, filename, duration = 50):
     imgs[0].save(filename, save_all = True, append_images = imgs[1:], duration = duration, loop = 0)
 
 def clean_filename(filename):
-    return filename.replace("_00000", "").replace("00000_", "")
+    return filename.replace("_000000", "").replace("000000_", "")
 
 # Save a list of images with ordering and according to a path template
 def save_images_builder(drange, ratio, grid_size, grid = False, verbose = False):

training/network.py

@@ -1,4 +1,4 @@
-# Network architectures for the GANSformer model, as well as multiple baselines such as
+# Network architectures for the GANformer model, as well as multiple baselines such as
 # vanilla GAN, StyleGAN2, k-GAN and SAGAN, all implemented as extensions of the same
 # model skeleton for most precise comparability under same settings.
 # See readme for architecture overview.
@@ -431,7 +431,7 @@ def modulated_conv2d_layer(x, y, dim, kernel, # Convolution parameters
 
     return x
 
-# Local attention block, potentially replacing convolution (not used by default in GANsformer)
+# Local attention block, potentially replacing convolution (not used by default in GANformer)
 # Grid is positional embeddings that match the size of a convolution receptive window
 # X shape is [NCHW]
 def att2d_layer(x, dim, kernel, grid, up = False, down = False, resample_kernel = None, num_heads = 4):
@@ -797,7 +797,7 @@ def merge_images(x, dlatents, k, type, same = False):
 
 # Generator network, composed of two sub-networks: mapping and synthesis, as defined below
 # This function is mostly similar to the original StyleGAN2 version
-def G_GANsformer(
+def G_GANformer(
     latents_in,                               # First input: Latent vectors (z) [batch_size, latent_size]
     labels_in,                                # Second input (optional): Conditioning labels [batch_size, label_size]
     # General arguments
@@ -810,7 +810,7 @@ def G_GANsformer(
     mapping_func            = "G_mapping",    # Function name of the mapping network
     synthesis_func          = "G_synthesis",  # Function name of the synthesis network
     # Truncation cutoff
-    truncation_psi          = 0.65,           # Style strength multiplier for the truncation trick. None = disable
+    truncation_psi          = 0.7,            # Style strength multiplier for the truncation trick. None = disable
     truncation_cutoff       = None,           # Number of layers for which to apply the truncation trick. None = disable
     truncation_psi_val      = None,           # Value for truncation_psi to use during validation
     truncation_cutoff_val   = None,           # Value for truncation_cutoff to use during validation
@@ -907,7 +907,7 @@ def G_GANsformer(
     # Mixing (see StyleGAN): mixes together some of the W latents mapped from one set of
     # source latents A, and another set of source latents B. Mixing can be between latents
     # that correspond either to different layers or to the different k components used in
-    # the GANsformer.
+    # the GANformer.
     # Used only during training, and by default only layers are mixes as in StyleGAN.
     def mixing(latents_in, dlatents, prob, num, idx):
         if prob is None or prob == 0:
@@ -960,7 +960,7 @@ def G_GANsformer(
 
 ########################################### Mapping network ###########################################
 
-# Mostly similar to the StylGAN version, with new options for the GANsformer such as:
+# Mostly similar to the StylGAN version, with new options for the GANformer such as:
 # self-attention among latents, shared mapping of components, and support for mapping_dim != output_dim
 def G_mapping(
     # Tensor inputs
@@ -977,7 +977,7 @@ def G_mapping(
     dlatent_broadcast       = None,         # Output disentangled latent (w) as [batch_size, dlatent_size]
                                             # or [batch_size, dlatent_broadcast, dlatent_size]
     normalize_latents       = True,         # Normalize latent vectors (z) before feeding them to the mapping layers?
-    transformer             = False,        # Whereas the generator uses attention or not (i.e. GANsformer or StyleGAN)
+    transformer             = False,        # Whereas the generator uses attention or not (i.e. GANformer or StyleGAN)
     # Mapping network
     mapping_layersnum       = 8,            # Number of mapping layers
     mapping_dim             = None,         # Number of activations in the mapping layers
@@ -1101,7 +1101,7 @@ def G_mapping(
 
 # Main differences from the StyleGAN version include the incorporation of transformer layers.
 # This function supports different generator forms:
-# - GANsformer (--transformer)
+# - GANformer (--transformer)
 # - GAN (--latent-stem --style = False)
 # - StyleGAN (--style)
 # - SAGAN (--img2img)
@@ -1120,7 +1120,7 @@ def G_synthesis(
     fmap_decay          = 1.0,          # log2 network dimension reduction when doubling the resolution
     fmap_min            = 1,            # Minimum network dimension in any layer
     fmap_max            = 512,          # Maximum network dimension in any layer
-    # General settings (StyleGAN/GANsformer related)
+    # General settings (StyleGAN/GANformer related)
     architecture        = "skip",       # Architecture: orig, skip, resnet
     nonlinearity        = "lrelu",      # Activation function: relu, lrelu", etc.
     resample_kernel     = [1, 3, 3, 1], # Low-pass filter to apply when resampling activations, None = no filtering
@@ -1131,8 +1131,8 @@ def G_synthesis(
     randomize_noise     = True,         # If True, randomize noise inputs every time (non-deterministic)
                                         # Otherwise, read noise inputs from variables
     tanh                = False,        # Use tanh on output activations (turns out to be unnecessary and suboptimal)
-    # GANsformer settings
-    transformer         = False,        # Whereas the generator uses attention or not (i.e. GANsformer or StyleGAN)
+    # GANformer settings
+    transformer         = False,        # Whereas the generator uses attention or not (i.e. GANformer or StyleGAN)
     components_num      = 1,            # Number of Latent (z) vector components z_1,...,z_k
     num_heads           = 1,            # Number of attention heads
     attention_dropout   = 0.12,         # Attention dropout rate
@@ -1206,7 +1206,7 @@ def G_synthesis(
     batch_size = get_shape(dlatents_in)[0]
 
     # Positional encodings (spatial)
-    # Image positional encodings, dictionary at different resolutions (used in GANsformer)
+    # Image positional encodings, dictionary at different resolutions (used in GANformer)
     grid_poses = get_positional_embeddings(resolution_log2, pos_dim, pos_type, pos_directions_num, pos_init)
     # Encodings for Local attention (not used by default)
     if local_attention:
@@ -1357,7 +1357,7 @@ def G_synthesis(
             return upsample_2d(y, k = resample_kernel)
 
     # Convert image features to output image (e.g. RGB)
-    # - For GANsformer: optionally perform transformer on highest resolution (if end_res >= resolution_log2)
+    # - For GANformer: optionally perform transformer on highest resolution (if end_res >= resolution_log2)
     # - For k-GAN: Merge the k images, if didn't merge so far
     # - optional tanh activation on outputs (not used by default)
     # - skip connections (for --g-arch skip)
@@ -1479,7 +1479,7 @@ def G_synthesis(
 # -----------------------------------------------------------------------------------------------------
 
 # Discriminator network, with convolution, downsampling, and optional transformer layers
-def D_GANsformer(
+def D_GANformer(
     images_in,                          # First input: Images [batch_size, channel, height, width]
     labels_in,                          # Second input: Labels [batch_size, label_size]
     # Dimensions and resolution
@@ -1498,7 +1498,7 @@ def D_GANsformer(
     mbstd_group_size    = 4,            # Group size for the minibatch standard deviation layer, 0 = disable
     mbstd_num_features  = 1,            # Number of features for the minibatch standard deviation layer
     resample_kernel     = [1, 3, 3, 1], # Low-pass filter to apply when resampling activations, None = no filtering
-    # GANsformer settings
+    # GANformer settings
     transformer         = False,        # Whereas the discriminator uses attention or not
     components_num      = 1,            # Number of aggregator variables
     num_heads           = 1,            # Number of attention heads

training/training_loop.py

@@ -216,12 +216,14 @@ def training_loop(
         misc.save_pkl((G, D, Gs), pkl, remove = False)
 
         # Quantitative evaluation
+        misc.log("Run evaluation...")
         metric = metrics.run(pkl, num_imgs = eval_images_num, run_dir = dnnlib.make_run_dir_path(),
             data_dir = dnnlib.convert_path(data_dir), num_gpus = num_gpus, ratio = ratio, 
-            tf_config = tf_config, mirror_augment = mirror_augment)  
+            tf_config = tf_config, eval_mod = True, mirror_augment = mirror_augment)  
 
         # Qualitative evaluation
-        visualize.eval(G, dataset, batch_size = sched.minibatch_gpu,
+        misc.log("Produce visualizations...")
+        visualize.eval(Gs, dataset, batch_size = sched.minibatch_gpu,
             drange_net = drange_net, ratio = ratio, **vis_args)
 
     if not train:
@@ -393,7 +395,7 @@ def training_loop(
 
             # Save snapshots
             if img_snapshot_ticks is not None and (cur_tick % img_snapshot_ticks == 0 or done):
-                visualize.eval(G, dataset, batch_size = sched.minibatch_gpu, training = True,
+                visualize.eval(Gs, dataset, batch_size = sched.minibatch_gpu, training = True,
                     step = cur_nimg // 1000, grid_size = grid_size, latents = grid_latents, 
                     labels = grid_labels, drange_net = drange_net, ratio = ratio, **vis_args)
 

training/visualize.py

@@ -43,8 +43,8 @@ def slerp(a, b, ts):
 ## Creates variety of visualizations. Supported types
 # imgs           : Save image samples
 # ltnts          : Save latent vectors
-# maps           : Save attention maps (for GANsformer only)
-# layer_maps     : Save attention maps for all layers (for GANsformer only)
+# maps           : Save attention maps (for GANformer only)
+# layer_maps     : Save attention maps for all layers (for GANformer only)
 # interpolations : Create latent interpolations
 # noise_var      : Create noise variation visualization
 # style_mix      : Create style mixing visualization
@@ -165,7 +165,7 @@ def eval(G,
             misc.save_npys(latents, pattern_of("latents-z", step, "npy"), verbose, idx)
             misc.save_npys(wlatents, pattern_of("latents-w", step, "npy"), verbose, idx)
 
-        # For the GANsformer model, save attention maps
+        # For the GANformer model, save attention maps
         if attention:
             if "maps" in vis:
                 pallete = np.expand_dims(misc.get_colors(components_num), axis = [2, 3])
@@ -209,7 +209,7 @@ def eval(G,
                     save_images(maps, pattern, idx)
 
     # Produce interpolations between pairs or source latents
-    # In the GANsformer case, varying one component at a time
+    # In the GANformer case, varying one component at a time
     if "interpolations" in vis:
         ts = np.array(np.linspace(0.0, 1.0, num = intrp_density, endpoint = True))
 
@@ -275,7 +275,7 @@ def eval(G,
         PIL.Image.fromarray(diff, "L").save(dnnlib.make_run_dir_path("eval/noise-variance.png"))
 
     # Compute style mixing table, varying using the latent A in some of the layers and latent B in rest.
-    # For the GANsformer, also produce component mixes (using latents from A in some of the components,
+    # For the GANformer, also produce component mixes (using latents from A in some of the components,
     # and latents from B in the rest.
     if "style_mix" in vis:
         if verbose: