Documentation

README.md

@@ -203,7 +203,7 @@ python dataset_tool.py --source=~/downloads/afhq/train/wild --dest=~/datasets/af
 python dataset_tool.py --source=~/downloads/cifar-10-python.tar.gz --dest=~/datasets/cifar10.zip
 ```
 
-**LSUN**: Download the desired LSUN categories in LMDB format from the [LSUN project page](https://www.yf.io/p/lsun/) and convert to ZIP archive:
+**LSUN**: Download the desired categories from the [LSUN project page](https://www.yf.io/p/lsun/) and convert to ZIP archive:
 
 ```.bash
 python dataset_tool.py --source=~/downloads/lsun/raw/cat_lmdb --dest=~/datasets/lsuncat200k.zip \
@@ -262,7 +262,7 @@ The training configuration can be further customized with additional command lin
 * `--cond=1` enables class-conditional training (requires a dataset with labels).
 * `--mirror=1` amplifies the dataset with x-flips. Often beneficial, even with ADA.
 * `--resume=ffhq1024 --snap=10` performs transfer learning from FFHQ trained at 1024x1024.
-* `--resume=~/training-runs/<NAME>/network-snapshot-<KIMG>.pkl` resumes a previous training run where it left off.
+* `--resume=~/training-runs/<NAME>/network-snapshot-<INT>.pkl` resumes a previous training run.
 * `--gamma=10` overrides R1 gamma. We recommend trying a couple of different values for each new dataset.
 * `--aug=ada --target=0.7` adjusts ADA target value (default: 0.6).
 * `--augpipe=blit` enables pixel blitting but disables all other augmentations.
@@ -293,7 +293,7 @@ The total training time depends heavily on resolution, number of GPUs, dataset,
 | 1024x1024  | 4    | 11h 36m   | 12d 02h    | 40.1&ndash;40.8   | 8.4 GB  | 21.9 GB
 | 1024x1024  | 8    | 5h 54m    | 6d 03h     | 20.2&ndash;20.6   | 8.3 GB  | 44.7 GB
 
-The above measurements were done using NVIDIA Tesla V100 GPUs with default settings (`--cfg=auto --aug=ada --metrics=fid50k_full`). "sec/kimg" shows the expected range of variation in raw training performance, as reported in `log.txt`, and "GPU mem" and "CPU mem" show the peak memory consumption observed over the course of training.
+The above measurements were done using NVIDIA Tesla V100 GPUs with default settings (`--cfg=auto --aug=ada --metrics=fid50k_full`). "sec/kimg" shows the expected range of variation in raw training performance, as reported in `log.txt`. "GPU mem" and "CPU mem" show the highest observed memory consumption, excluding the peak at the beginning caused by `torch.backends.cudnn.benchmark`.
 
 In typical cases, 25000 kimg or more is needed to reach convergence, but the results are already quite reasonable around 5000 kimg. 1000 kimg is often enough for transfer learning, which tends to converge significantly faster. The following figure shows example convergence curves for different datasets as a function of wallclock time, using the same settings as above:
 
@@ -325,23 +325,23 @@ We employ the following metrics in the ADA paper. Execution time and GPU memory
 
 | Metric        | Time   | GPU mem | Description |
 | :-----        | :----: | :-----: | :---------- |
-| `fid50k_full` | 13 min | 1.8 GB  | Fr&eacute;chet inception distance<sup>[1]</sup> against the full dataset.
-| `kid50k_full` | 13 min | 1.8 GB  | Kernel inception distance<sup>[2]</sup> against the full dataset.
-| `pr50k3_full` | 13 min | 4.1 GB  | Precision and recall<sup>[3]</sup> againt the full dataset.
-| `is50k`       | 13 min | 1.8 GB  | Inception score<sup>[4]</sup> for CIFAR-10.
+| `fid50k_full` | 13 min | 1.8 GB  | Fr&eacute;chet inception distance<sup>[1]</sup> against the full dataset
+| `kid50k_full` | 13 min | 1.8 GB  | Kernel inception distance<sup>[2]</sup> against the full dataset
+| `pr50k3_full` | 13 min | 4.1 GB  | Precision and recall<sup>[3]</sup> againt the full dataset
+| `is50k`       | 13 min | 1.8 GB  | Inception score<sup>[4]</sup> for CIFAR-10
 
 In addition, the following metrics from the [StyleGAN](https://github.com/NVlabs/stylegan) and [StyleGAN2](https://github.com/NVlabs/stylegan2) papers are also supported:
 
 | Metric        | Time   | GPU mem | Description |
 | :------------ | :----: | :-----: | :---------- |
-| `fid50k`      | 13 min | 1.8 GB  | Fr&eacute;chet inception distance against 50k real images.
-| `kid50k`      | 13 min | 1.8 GB  | Kernel inception distance against 50k real images.
-| `pr50k3`      | 13 min | 4.1 GB  | Precision and recall against 50k real images.
-| `ppl2_wend`   | 36 min | 2.4 GB  | Perceptual path length<sup>[5]</sup> in W at path endpoints against full image.
-| `ppl_zfull`   | 36 min | 2.4 GB  | Perceptual path length in Z for full paths against cropped image.
-| `ppl_wfull`   | 36 min | 2.4 GB  | Perceptual path length in W for full paths against cropped image.
-| `ppl_zend`    | 36 min | 2.4 GB  | Perceptual path length in Z at path endpoints against cropped image.
-| `ppl_wend`    | 36 min | 2.4 GB  | Perceptual path length in W at path endpoints against cropped image.
+| `fid50k`      | 13 min | 1.8 GB  | Fr&eacute;chet inception distance against 50k real images
+| `kid50k`      | 13 min | 1.8 GB  | Kernel inception distance against 50k real images
+| `pr50k3`      | 13 min | 4.1 GB  | Precision and recall against 50k real images
+| `ppl2_wend`   | 36 min | 2.4 GB  | Perceptual path length<sup>[5]</sup> in W, endpoints, full image
+| `ppl_zfull`   | 36 min | 2.4 GB  | Perceptual path length in Z, full paths, cropped image
+| `ppl_wfull`   | 36 min | 2.4 GB  | Perceptual path length in W, full paths, cropped image
+| `ppl_zend`    | 36 min | 2.4 GB  | Perceptual path length in Z, endpoints, cropped image
+| `ppl_wend`    | 36 min | 2.4 GB  | Perceptual path length in W, endpoints, cropped image
 
 References:
 1. [GANs Trained by a Two Time-Scale Update Rule Converge to a Local Nash Equilibrium](https://arxiv.org/abs/1706.08500), Heusel et al. 2017

dnnlib/__init__.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
 #
 # NVIDIA CORPORATION and its licensors retain all intellectual property
 # and proprietary rights in and to this software, related documentation

docs/dataset-tool-help.txt

@@ -1,50 +1,50 @@
-Usage: dataset_tool.py [OPTIONS]
-
-  Convert an image dataset into a dataset archive usable with StyleGAN2 ADA
-  PyTorch.
-
-  The input dataset format is guessed from the --source argument:
-
-  --source *_lmdb/                - Load LSUN dataset
-  --source cifar-10-python.tar.gz - Load CIFAR-10 dataset
-  --source path/                  - Recursively load all images from path/
-  --source dataset.zip            - Recursively load all images from dataset.zip
-
-  The output dataset format can be either an image folder or a zip archive.
-  Specifying the output format and path:
-
-  --dest /path/to/dir             - Save output files under /path/to/dir
-  --dest /path/to/dataset.zip     - Save output files into /path/to/dataset.zip archive
-
-  Images within the dataset archive will be stored as uncompressed PNG.
-
-  Image scale/crop and resolution requirements:
-
-  Output images must be square-shaped and they must all have the same power-
-  of-two dimensions.
-
-  To scale arbitrary input image size to a specific width and height, use
-  the --width and --height options.  Output resolution will be either the
-  original input resolution (if --width/--height was not specified) or the
-  one specified with --width/height.
-
-  Use the --transform=center-crop or --transform=center-crop-wide options to
-  apply a center crop transform on the input image.  These options should be
-  used with the --width and --height options.  For example:
-
-  python dataset_tool.py --source LSUN/raw/cat_lmdb --dest /tmp/lsun_cat \
-      --transform=center-crop-wide --width 512 --height=384
-
-Options:
-  --source PATH                   Directory or archive name for input dataset
-                                  [required]
-  --dest PATH                     Output directory or archive name for output
-                                  dataset  [required]
-  --max-images INTEGER            Output only up to `max-images` images
-  --resize-filter [box|lanczos]   Filter to use when resizing images for
-                                  output resolution  [default: lanczos]
-  --transform [center-crop|center-crop-wide]
-                                  Input crop/resize mode
-  --width INTEGER                 Output width
-  --height INTEGER                Output height
-  --help                          Show this message and exit.
+Usage: dataset_tool.py [OPTIONS]
+
+  Convert an image dataset into a dataset archive usable with StyleGAN2 ADA
+  PyTorch.
+
+  The input dataset format is guessed from the --source argument:
+
+  --source *_lmdb/                - Load LSUN dataset
+  --source cifar-10-python.tar.gz - Load CIFAR-10 dataset
+  --source path/                  - Recursively load all images from path/
+  --source dataset.zip            - Recursively load all images from dataset.zip
+
+  The output dataset format can be either an image folder or a zip archive.
+  Specifying the output format and path:
+
+  --dest /path/to/dir             - Save output files under /path/to/dir
+  --dest /path/to/dataset.zip     - Save output files into /path/to/dataset.zip archive
+
+  Images within the dataset archive will be stored as uncompressed PNG.
+
+  Image scale/crop and resolution requirements:
+
+  Output images must be square-shaped and they must all have the same power-
+  of-two dimensions.
+
+  To scale arbitrary input image size to a specific width and height, use
+  the --width and --height options.  Output resolution will be either the
+  original input resolution (if --width/--height was not specified) or the
+  one specified with --width/height.
+
+  Use the --transform=center-crop or --transform=center-crop-wide options to
+  apply a center crop transform on the input image.  These options should be
+  used with the --width and --height options.  For example:
+
+  python dataset_tool.py --source LSUN/raw/cat_lmdb --dest /tmp/lsun_cat \
+      --transform=center-crop-wide --width 512 --height=384
+
+Options:
+  --source PATH                   Directory or archive name for input dataset
+                                  [required]
+  --dest PATH                     Output directory or archive name for output
+                                  dataset  [required]
+  --max-images INTEGER            Output only up to `max-images` images
+  --resize-filter [box|lanczos]   Filter to use when resizing images for
+                                  output resolution  [default: lanczos]
+  --transform [center-crop|center-crop-wide]
+                                  Input crop/resize mode
+  --width INTEGER                 Output width
+  --height INTEGER                Output height
+  --help                          Show this message and exit.

docs/train-help.txt

@@ -1,69 +1,69 @@
-Usage: train.py [OPTIONS]
-
-  Train a GAN using the techniques described in the paper "Training
-  Generative Adversarial Networks with Limited Data".
-
-  Examples:
-
-  # Train with custom images using 1 GPU.
-  python train.py --outdir=~/training-runs --data=~/my-image-folder
-
-  # Train class-conditional CIFAR-10 using 2 GPUs.
-  python train.py --outdir=~/training-runs --data=~/datasets/cifar10.zip \
-      --gpus=2 --cfg=cifar --cond=1
-
-  # Transfer learn MetFaces from FFHQ using 4 GPUs.
-  python train.py --outdir=~/training-runs --data=~/datasets/metfaces.zip \
-      --gpus=4 --cfg=paper1024 --mirror=1 --resume=ffhq1024 --snap=10
-
-  # Reproduce original StyleGAN2 config F.
-  python train.py --outdir=~/training-runs --data=~/datasets/ffhq.zip \
-      --gpus=8 --cfg=stylegan2 --mirror=1 --aug=noaug
-
-  Base configs (--cfg):
-    auto       Automatically select reasonable defaults based on resolution
-               and GPU count. Good starting point for new datasets.
-    stylegan2  Reproduce results for StyleGAN2 config F at 1024x1024.
-    paper256   Reproduce results for FFHQ and LSUN Cat at 256x256.
-    paper512   Reproduce results for BreCaHAD and AFHQ at 512x512.
-    paper1024  Reproduce results for MetFaces at 1024x1024.
-    cifar      Reproduce results for CIFAR-10 at 32x32.
-
-  Transfer learning source networks (--resume):
-    ffhq256        FFHQ trained at 256x256 resolution.
-    ffhq512        FFHQ trained at 512x512 resolution.
-    ffhq1024       FFHQ trained at 1024x1024 resolution.
-    celebahq256    CelebA-HQ trained at 256x256 resolution.
-    lsundog256     LSUN Dog trained at 256x256 resolution.
-    <PATH or URL>  Custom network pickle.
-
-Options:
-  --outdir DIR                    Where to save the results  [required]
-  --gpus INT                      Number of GPUs to use [default: 1]
-  --snap INT                      Snapshot interval [default: 50 ticks]
-  --metrics LIST                  Comma-separated list or "none" [default:
-                                  fid50k_full]
-  --seed INT                      Random seed [default: 0]
-  -n, --dry-run                   Print training options and exit
-  --data PATH                     Training data (directory or zip)  [required]
-  --cond BOOL                     Train conditional model based on dataset
-                                  labels [default: false]
-  --subset INT                    Train with only N images [default: all]
-  --mirror BOOL                   Enable dataset x-flips [default: false]
-  --cfg [auto|stylegan2|paper256|paper512|paper1024|cifar]
-                                  Base config [default: auto]
-  --gamma FLOAT                   Override R1 gamma
-  --kimg INT                      Override training duration
-  --batch INT                     Override batch size
-  --aug [noaug|ada|fixed]         Augmentation mode [default: ada]
-  --p FLOAT                       Augmentation probability for --aug=fixed
-  --target FLOAT                  ADA target value for --aug=ada
-  --augpipe [blit|geom|color|filter|noise|cutout|bg|bgc|bgcf|bgcfn|bgcfnc]
-                                  Augmentation pipeline [default: bgc]
-  --resume PKL                    Resume training [default: noresume]
-  --freezed INT                   Freeze-D [default: 0 layers]
-  --fp32 BOOL                     Disable mixed-precision training
-  --nhwc BOOL                     Use NHWC memory format with FP16
-  --nobench BOOL                  Disable cuDNN benchmarking
-  --workers INT                   Override number of DataLoader workers
-  --help                          Show this message and exit.
+Usage: train.py [OPTIONS]
+
+  Train a GAN using the techniques described in the paper "Training
+  Generative Adversarial Networks with Limited Data".
+
+  Examples:
+
+  # Train with custom images using 1 GPU.
+  python train.py --outdir=~/training-runs --data=~/my-image-folder
+
+  # Train class-conditional CIFAR-10 using 2 GPUs.
+  python train.py --outdir=~/training-runs --data=~/datasets/cifar10.zip \
+      --gpus=2 --cfg=cifar --cond=1
+
+  # Transfer learn MetFaces from FFHQ using 4 GPUs.
+  python train.py --outdir=~/training-runs --data=~/datasets/metfaces.zip \
+      --gpus=4 --cfg=paper1024 --mirror=1 --resume=ffhq1024 --snap=10
+
+  # Reproduce original StyleGAN2 config F.
+  python train.py --outdir=~/training-runs --data=~/datasets/ffhq.zip \
+      --gpus=8 --cfg=stylegan2 --mirror=1 --aug=noaug
+
+  Base configs (--cfg):
+    auto       Automatically select reasonable defaults based on resolution
+               and GPU count. Good starting point for new datasets.
+    stylegan2  Reproduce results for StyleGAN2 config F at 1024x1024.
+    paper256   Reproduce results for FFHQ and LSUN Cat at 256x256.
+    paper512   Reproduce results for BreCaHAD and AFHQ at 512x512.
+    paper1024  Reproduce results for MetFaces at 1024x1024.
+    cifar      Reproduce results for CIFAR-10 at 32x32.
+
+  Transfer learning source networks (--resume):
+    ffhq256        FFHQ trained at 256x256 resolution.
+    ffhq512        FFHQ trained at 512x512 resolution.
+    ffhq1024       FFHQ trained at 1024x1024 resolution.
+    celebahq256    CelebA-HQ trained at 256x256 resolution.
+    lsundog256     LSUN Dog trained at 256x256 resolution.
+    <PATH or URL>  Custom network pickle.
+
+Options:
+  --outdir DIR                    Where to save the results  [required]
+  --gpus INT                      Number of GPUs to use [default: 1]
+  --snap INT                      Snapshot interval [default: 50 ticks]
+  --metrics LIST                  Comma-separated list or "none" [default:
+                                  fid50k_full]
+  --seed INT                      Random seed [default: 0]
+  -n, --dry-run                   Print training options and exit
+  --data PATH                     Training data (directory or zip)  [required]
+  --cond BOOL                     Train conditional model based on dataset
+                                  labels [default: false]
+  --subset INT                    Train with only N images [default: all]
+  --mirror BOOL                   Enable dataset x-flips [default: false]
+  --cfg [auto|stylegan2|paper256|paper512|paper1024|cifar]
+                                  Base config [default: auto]
+  --gamma FLOAT                   Override R1 gamma
+  --kimg INT                      Override training duration
+  --batch INT                     Override batch size
+  --aug [noaug|ada|fixed]         Augmentation mode [default: ada]
+  --p FLOAT                       Augmentation probability for --aug=fixed
+  --target FLOAT                  ADA target value for --aug=ada
+  --augpipe [blit|geom|color|filter|noise|cutout|bg|bgc|bgcf|bgcfn|bgcfnc]
+                                  Augmentation pipeline [default: bgc]
+  --resume PKL                    Resume training [default: noresume]
+  --freezed INT                   Freeze-D [default: 0 layers]
+  --fp32 BOOL                     Disable mixed-precision training
+  --nhwc BOOL                     Use NHWC memory format with FP16
+  --nobench BOOL                  Disable cuDNN benchmarking
+  --workers INT                   Override number of DataLoader workers
+  --help                          Show this message and exit.

metrics/frechet_inception_distance.py

@@ -6,16 +6,21 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
+"""Frechet Inception Distance (FID) from the paper
+"GANs trained by a two time-scale update rule converge to a local Nash
+equilibrium". Matches the original implementation by Heusel et al. at
+https://github.com/bioinf-jku/TTUR/blob/master/fid.py"""
+
 import numpy as np
 import scipy.linalg
-
 from . import metric_utils
 
 #----------------------------------------------------------------------------
 
 def compute_fid(opts, max_real, num_gen):
+    # Direct TorchScript translation of http://download.tensorflow.org/models/image/imagenet/inception-2015-12-05.tgz
     detector_url = 'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metrics/inception-2015-12-05.pt'
-    detector_kwargs = dict(return_features=True)
+    detector_kwargs = dict(return_features=True) # Return raw features before the softmax layer.
 
     mu_real, sigma_real = metric_utils.compute_feature_stats_for_dataset(
         opts=opts, detector_url=detector_url, detector_kwargs=detector_kwargs,

metrics/inception_score.py

@@ -6,15 +6,19 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
-import numpy as np
+"""Inception Score (IS) from the paper "Improved techniques for training
+GANs". Matches the original implementation by Salimans et al. at
+https://github.com/openai/improved-gan/blob/master/inception_score/model.py"""
 
+import numpy as np
 from . import metric_utils
 
 #----------------------------------------------------------------------------
 
 def compute_is(opts, num_gen, num_splits):
+    # Direct TorchScript translation of http://download.tensorflow.org/models/image/imagenet/inception-2015-12-05.tgz
     detector_url = 'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metrics/inception-2015-12-05.pt'
-    detector_kwargs = dict(no_output_bias=True)
+    detector_kwargs = dict(no_output_bias=True) # Match the original implementation by not applying bias in the softmax layer.
 
     gen_probs = metric_utils.compute_feature_stats_for_generator(
         opts=opts, detector_url=detector_url, detector_kwargs=detector_kwargs,

metrics/kernel_inception_distance.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
 #
 # NVIDIA CORPORATION and its licensors retain all intellectual property
 # and proprietary rights in and to this software, related documentation
@@ -6,15 +6,19 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
-import numpy as np
+"""Kernel Inception Distance (KID) from the paper "Demystifying MMD
+GANs". Matches the original implementation by Binkowski et al. at
+https://github.com/mbinkowski/MMD-GAN/blob/master/gan/compute_scores.py"""
 
+import numpy as np
 from . import metric_utils
 
 #----------------------------------------------------------------------------
 
 def compute_kid(opts, max_real, num_gen, num_subsets, max_subset_size):
+    # Direct TorchScript translation of http://download.tensorflow.org/models/image/imagenet/inception-2015-12-05.tgz
     detector_url = 'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metrics/inception-2015-12-05.pt'
-    detector_kwargs = dict(return_features=True)
+    detector_kwargs = dict(return_features=True) # Return raw features before the softmax layer.
 
     real_features = metric_utils.compute_feature_stats_for_dataset(
         opts=opts, detector_url=detector_url, detector_kwargs=detector_kwargs,

metrics/perceptual_path_length.py

@@ -6,11 +6,15 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
+"""Perceptual Path Length (PPL) from the paper "A Style-Based Generator
+Architecture for Generative Adversarial Networks". Matches the original
+implementation by Karras et al. at
+https://github.com/NVlabs/stylegan/blob/master/metrics/perceptual_path_length.py"""
+
 import copy
 import numpy as np
 import torch
 import dnnlib
-
 from . import metric_utils
 
 #----------------------------------------------------------------------------

metrics/precision_recall.py

@@ -6,8 +6,12 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
-import torch
+"""Precision/Recall (PR) from the paper "Improved Precision and Recall
+Metric for Assessing Generative Models". Matches the original implementation
+by Kynkaanniemi et al. at
+https://github.com/kynkaat/improved-precision-and-recall-metric/blob/master/precision_recall.py"""
 
+import torch
 from . import metric_utils
 
 #----------------------------------------------------------------------------

torch_utils/persistence.py

@@ -6,6 +6,13 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
+"""Facilities for pickling Python code alongside other data.
+
+The pickled code is automatically imported into a separate Python module
+during unpickling. This way, any previously exported pickles will remain
+usable even if the original code is no longer available, or if the current
+version of the code is not consistent with what was originally pickled."""
+
 import sys
 import pickle
 import io
@@ -17,29 +24,70 @@ import dnnlib
 
 #----------------------------------------------------------------------------
 
-_version = 6
-_decorators = set() # {decorator_class}
-_import_hooks = [] # [function]
-_module_to_src_dict = dict() # {module: src}
-_src_to_module_dict = dict() # {src: module}
-
-#----------------------------------------------------------------------------
-
-def is_persistent(obj):
-    try:
-        if obj in _decorators:
-            return True
-    except TypeError:
-        pass
-    return type(obj) in _decorators # pylint: disable=unidiomatic-typecheck
-
-def import_hook(func):
-    assert callable(func)
-    _import_hooks.append(func)
+_version            = 6         # internal version number
+_decorators         = set()     # {decorator_class, ...}
+_import_hooks       = []        # [hook_function, ...]
+_module_to_src_dict = dict()    # {module: src, ...}
+_src_to_module_dict = dict()    # {src: module, ...}
 
 #----------------------------------------------------------------------------
 
 def persistent_class(orig_class):
+    r"""Class decorator that extends a given class to save its source code
+    when pickled.
+
+    Example:
+
+        from torch_utils import persistence
+
+        @persistence.persistent_class
+        class MyNetwork(torch.nn.Module):
+            def __init__(self, num_inputs, num_outputs):
+                super().__init__()
+                self.fc = MyLayer(num_inputs, num_outputs)
+                ...
+
+        @persistence.persistent_class
+        class MyLayer(torch.nn.Module):
+            ...
+
+    When pickled, any instance of `MyNetwork` and `MyLayer` will save its
+    source code alongside other internal state (e.g., parameters, buffers,
+    and submodules). This way, any previously exported pickle will remain
+    usable even if the class definitions have been modified or are no
+    longer available.
+
+    The decorator saves the source code of the entire Python module
+    containing the decorated class. It does *not* save the source code of
+    any imported modules. Thus, the imported modules must be available
+    during unpickling, also including `torch_utils.persistence` itself.
+
+    It is ok to call functions defined in the same module from the
+    decorated class. However, if the decorated class depends on other
+    classes defined in the same module, they must be decorated as well.
+    This is illustrated in the above example in the case of `MyLayer`.
+
+    It is also possible to employ the decorator just-in-time before
+    calling the constructor. For example:
+
+        cls = MyLayer
+        if want_to_make_it_persistent:
+            cls = persistence.persistent_class(cls)
+        layer = cls(num_inputs, num_outputs)
+
+    As an additional feature, the decorator also keeps track of the
+    arguments that were used to construct each instance of the decorated
+    class. The arguments can be queried via `obj.init_args` and
+    `obj.init_kwargs`, and they are automatically pickled alongside other
+    object state. A typical use case is to first unpickle a previous
+    instance of a persistent class, and then upgrade it to use the latest
+    version of the source code:
+
+        with open('old_pickle.pkl', 'rb') as f:
+            old_net = pickle.load(f)
+        new_net = MyNetwork(*old_obj.init_args, **old_obj.init_kwargs)
+        misc.copy_params_and_buffers(old_net, new_net, require_all=True)
+    """
     assert isinstance(orig_class, type)
     if is_persistent(orig_class):
         return orig_class
@@ -83,7 +131,55 @@ def persistent_class(orig_class):
 
 #----------------------------------------------------------------------------
 
+def is_persistent(obj):
+    r"""Test whether the given object or class is persistent, i.e.,
+    whether it will save its source code when pickled.
+    """
+    try:
+        if obj in _decorators:
+            return True
+    except TypeError:
+        pass
+    return type(obj) in _decorators # pylint: disable=unidiomatic-typecheck
+
+#----------------------------------------------------------------------------
+
+def import_hook(hook):
+    r"""Register an import hook that is called whenever a persistent object
+    is being unpickled. A typical use case is to patch the pickled source
+    code to avoid errors and inconsistencies when the API of some imported
+    module has changed.
+
+    The hook should have the following signature:
+
+        hook(meta) -> modified meta
+
+    `meta` is an instance of `dnnlib.EasyDict` with the following fields:
+
+        type:       Type of the persistent object, e.g. `'class'`.
+        version:    Internal version number of `torch_utils.persistence`.
+        module_src  Original source code of the Python module.
+        class_name: Class name in the original Python module.
+        state:      Internal state of the object.
+
+    Example:
+
+        @persistence.import_hook
+        def wreck_my_network(meta):
+            if meta.class_name == 'MyNetwork':
+                print('MyNetwork is being imported. I will wreck it!')
+                meta.module_src = meta.module_src.replace("True", "False")
+            return meta
+    """
+    assert callable(hook)
+    _import_hooks.append(hook)
+
+#----------------------------------------------------------------------------
+
 def _reconstruct_persistent_obj(meta):
+    r"""Hook that is called internally by the `pickle` module to unpickle
+    a persistent object.
+    """
     meta = dnnlib.EasyDict(meta)
     meta.state = dnnlib.EasyDict(meta.state)
     for hook in _import_hooks:
@@ -108,6 +204,8 @@ def _reconstruct_persistent_obj(meta):
 #----------------------------------------------------------------------------
 
 def _module_to_src(module):
+    r"""Query the source code of a given Python module.
+    """
     src = _module_to_src_dict.get(module, None)
     if src is None:
         src = inspect.getsource(module)
@@ -116,6 +214,8 @@ def _module_to_src(module):
     return src
 
 def _src_to_module(src):
+    r"""Get or create a Python module for the given source code.
+    """
     module = _src_to_module_dict.get(src, None)
     if module is None:
         module_name = "_imported_module_" + uuid.uuid4().hex
@@ -129,15 +229,19 @@ def _src_to_module(src):
 #----------------------------------------------------------------------------
 
 def _check_pickleable(obj):
+    r"""Check that the given object is pickleable, raising an exception if
+    it is not. This function is expected to be considerably more efficient
+    than actually pickling the object.
+    """
     def recurse(obj):
         if isinstance(obj, (list, tuple, set)):
             return [recurse(x) for x in obj]
         if isinstance(obj, dict):
             return [[recurse(x), recurse(y)] for x, y in obj.items()]
         if isinstance(obj, (str, int, float, bool, bytes, bytearray)):
-            return None # Primitive types are pickleable.
+            return None # Python primitive types are pickleable.
         if f'{type(obj).__module__}.{type(obj).__name__}' in ['numpy.ndarray', 'torch.Tensor']:
-            return None # Tensors are pickleable.
+            return None # NumPy arrays and PyTorch tensors are pickleable.
         if is_persistent(obj):
             return None # Persistent objects are pickleable, by virtue of the constructor check.
         return obj

torch_utils/training_stats.py

@@ -6,6 +6,11 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
+"""Facilities for reporting and collecting training statistics across
+multiple processes and devices. The interface is designed to minimize
+synchronization overhead as well as the amount of boilerplate in user
+code."""
+
 import re
 import numpy as np
 import torch
@@ -15,19 +20,31 @@ from . import misc
 
 #----------------------------------------------------------------------------
 
-_num_moments    = 3             # [num_scalars, sum_scalars, sum_squares]
+_num_moments    = 3             # [num_scalars, sum_of_scalars, sum_of_squares]
 _reduce_dtype   = torch.float32 # Data type to use for initial per-tensor reduction.
-_counter_dtype  = torch.float64 # Data type to use for the counters.
-
+_counter_dtype  = torch.float64 # Data type to use for the internal counters.
 _rank           = 0             # Rank of the current process.
 _sync_device    = None          # Device to use for multiprocess communication. None = single-process.
 _sync_called    = False         # Has _sync() been called yet?
-_counters       = dict()        # Running counter on each device, updated by report(): name => device => torch.Tensor
-_cumulative     = dict()        # Cumulative counter on the CPU, updated by _sync(): name => torch.Tensor
+_counters       = dict()        # Running counters on each device, updated by report(): name => device => torch.Tensor
+_cumulative     = dict()        # Cumulative counters on the CPU, updated by _sync(): name => torch.Tensor
 
 #----------------------------------------------------------------------------
 
 def init_multiprocessing(rank, sync_device):
+    r"""Initializes `torch_utils.training_stats` for collecting statistics
+    across multiple processes.
+
+    This function must be called after
+    `torch.distributed.init_process_group()` and before `Collector.update()`.
+    The call is not necessary if multi-process collection is not needed.
+
+    Args:
+        rank:           Rank of the current process.
+        sync_device:    PyTorch device to use for inter-process
+                        communication, or None to disable multi-process
+                        collection. Typically `torch.device('cuda', rank)`.
+    """
     global _rank, _sync_device
     assert not _sync_called
     _rank = rank
@@ -37,6 +54,28 @@ def init_multiprocessing(rank, sync_device):
 
 @misc.profiled_function
 def report(name, value):
+    r"""Broadcasts the given set of scalars to all interested instances of
+    `Collector`, across device and process boundaries.
+
+    This function is expected to be extremely cheap and can be safely
+    called from anywhere in the training loop, loss function, or inside a
+    `torch.nn.Module`.
+
+    Warning: The current implementation expects the set of unique names to
+    be consistent across processes. Please make sure that `report()` is
+    called at least once for each unique name by each process, and in the
+    same order. If a given process has no scalars to broadcast, it can do
+    `report(name, [])` (empty list).
+
+    Args:
+        name:   Arbitrary string specifying the name of the statistic.
+                Averages are accumulated separately for each unique name.
+        value:  Arbitrary set of scalars. Can be a list, tuple,
+                NumPy array, PyTorch tensor, or Python scalar.
+
+    Returns:
+        The same `value` that was passed in.
+    """
     if name not in _counters:
         _counters[name] = dict()
 
@@ -45,7 +84,11 @@ def report(name, value):
         return value
 
     elems = elems.detach().flatten().to(_reduce_dtype)
-    moments = torch.stack([torch.ones_like(elems).sum(), elems.sum(), elems.square().sum()])
+    moments = torch.stack([
+        torch.ones_like(elems).sum(),
+        elems.sum(),
+        elems.square().sum(),
+    ])
     assert moments.ndim == 1 and moments.shape[0] == _num_moments
     moments = moments.to(_counter_dtype)
 
@@ -58,45 +101,35 @@ def report(name, value):
 #----------------------------------------------------------------------------
 
 def report0(name, value):
+    r"""Broadcasts the given set of scalars by the first process (`rank = 0`),
+    but ignores any scalars provided by the other processes.
+    See `report()` for further details.
+    """
     report(name, value if _rank == 0 else [])
     return value
 
 #----------------------------------------------------------------------------
 
-def _sync(names):
-    if len(names) == 0:
-        return []
-    global _sync_called
-    _sync_called = True
-
-    # Collect deltas within current rank.
-    deltas = []
-    device = _sync_device if _sync_device is not None else torch.device('cpu')
-    for name in names:
-        delta = torch.zeros([_num_moments], dtype=_counter_dtype, device=device)
-        for counter in _counters[name].values():
-            delta.add_(counter.to(device))
-            counter.copy_(torch.zeros_like(counter))
-        deltas.append(delta)
-    deltas = torch.stack(deltas)
-
-    # Sum deltas across ranks.
-    if _sync_device is not None:
-        torch.distributed.all_reduce(deltas)
-
-    # Update cumulative values.
-    deltas = deltas.cpu()
-    for idx, name in enumerate(names):
-        if name not in _cumulative:
-            _cumulative[name] = torch.zeros([_num_moments], dtype=_counter_dtype)
-        _cumulative[name].add_(deltas[idx])
+class Collector:
+    r"""Collects the scalars broadcasted by `report()` and `report0()` and
+    computes their long-term averages (mean and standard deviation) over
+    user-defined periods of time.
 
-    # Return name-value pairs.
-    return [(name, _cumulative[name]) for name in names]
+    The averages are first collected into internal counters that are not
+    directly visible to the user. They are then copied to the user-visible
+    state as a result of calling `update()` and can then be queried using
+    `mean()`, `std()`, `as_dict()`, etc. Calling `update()` also resets the
+    internal counters for the next round, so that the user-visible state
+    effectively reflects averages collected between the last two calls to
+    `update()`.
 
-#----------------------------------------------------------------------------
-
-class Collector:
+    Args:
+        regex:          Regular expression defining which statistics to
+                        collect. The default is to collect everything.
+        keep_previous:  Whether to retain the previous averages if no
+                        scalars were collected on a given round
+                        (default: True).
+    """
     def __init__(self, regex='.*', keep_previous=True):
         self._regex = re.compile(regex)
         self._keep_previous = keep_previous
@@ -106,9 +139,24 @@ class Collector:
         self._moments.clear()
 
     def names(self):
+        r"""Returns the names of all statistics broadcasted so far that
+        match the regular expression specified at construction time.
+        """
         return [name for name in _counters if self._regex.fullmatch(name)]
 
     def update(self):
+        r"""Copies current values of the internal counters to the
+        user-visible state and resets them for the next round.
+
+        If `keep_previous=True` was specified at construction time, the
+        operation is skipped for statistics that have received no scalars
+        since the last update, retaining their previous averages.
+
+        This method performs a number of GPU-to-CPU transfers and one
+        `torch.distributed.all_reduce()`. It is intended to be called
+        periodically in the main training loop, typically once every
+        N training steps.
+        """
         if not self._keep_previous:
             self._moments.clear()
         for name, cumulative in _sync(self.names()):
@@ -120,22 +168,38 @@ class Collector:
                 self._moments[name] = delta
 
     def _get_delta(self, name):
+        r"""Returns the raw moments that were accumulated for the given
+        statistic between the last two calls to `update()`, or zero if
+        no scalars were collected.
+        """
         assert self._regex.fullmatch(name)
         if name not in self._moments:
             self._moments[name] = torch.zeros([_num_moments], dtype=_counter_dtype)
         return self._moments[name]
 
     def num(self, name):
+        r"""Returns the number of scalars that were accumulated for the given
+        statistic between the last two calls to `update()`, or zero if
+        no scalars were collected.
+        """
         delta = self._get_delta(name)
         return int(delta[0])
 
     def mean(self, name):
+        r"""Returns the mean of the scalars that were accumulated for the
+        given statistic between the last two calls to `update()`, or NaN if
+        no scalars were collected.
+        """
         delta = self._get_delta(name)
         if int(delta[0]) == 0:
             return float('nan')
         return float(delta[1] / delta[0])
 
     def std(self, name):
+        r"""Returns the standard deviation of the scalars that were
+        accumulated for the given statistic between the last two calls to
+        `update()`, or NaN if no scalars were collected.
+        """
         delta = self._get_delta(name)
         if int(delta[0]) == 0 or not np.isfinite(float(delta[1])):
             return float('nan')
@@ -146,12 +210,59 @@ class Collector:
         return np.sqrt(max(raw_var - np.square(mean), 0))
 
     def as_dict(self):
+        r"""Returns the averages accumulated between the last two calls to
+        `update()` as an `dnnlib.EasyDict`. The contents are as follows:
+
+            dnnlib.EasyDict(
+                NAME = dnnlib.EasyDict(num=FLOAT, mean=FLOAT, std=FLOAT),
+                ...
+            )
+        """
         stats = dnnlib.EasyDict()
         for name in self.names():
             stats[name] = dnnlib.EasyDict(num=self.num(name), mean=self.mean(name), std=self.std(name))
         return stats
 
     def __getitem__(self, name):
+        r"""Convenience getter.
+        `collector[name]` is a synonym for `collector.mean(name)`.
+        """
         return self.mean(name)
 
 #----------------------------------------------------------------------------
+
+def _sync(names):
+    r"""Synchronize the global cumulative counters across devices and
+    processes. Called internally by `Collector.update()`.
+    """
+    if len(names) == 0:
+        return []
+    global _sync_called
+    _sync_called = True
+
+    # Collect deltas within current rank.
+    deltas = []
+    device = _sync_device if _sync_device is not None else torch.device('cpu')
+    for name in names:
+        delta = torch.zeros([_num_moments], dtype=_counter_dtype, device=device)
+        for counter in _counters[name].values():
+            delta.add_(counter.to(device))
+            counter.copy_(torch.zeros_like(counter))
+        deltas.append(delta)
+    deltas = torch.stack(deltas)
+
+    # Sum deltas across ranks.
+    if _sync_device is not None:
+        torch.distributed.all_reduce(deltas)
+
+    # Update cumulative values.
+    deltas = deltas.cpu()
+    for idx, name in enumerate(names):
+        if name not in _cumulative:
+            _cumulative[name] = torch.zeros([_num_moments], dtype=_counter_dtype)
+        _cumulative[name].add_(deltas[idx])
+
+    # Return name-value pairs.
+    return [(name, _cumulative[name]) for name in names]
+
+#----------------------------------------------------------------------------