add files

compute_direction.py

@@ -0,0 +1,96 @@
+# python3.7
+"""Computes the semantic directions regarding a specific image region."""
+
+import os
+import argparse
+import numpy as np
+from tqdm import tqdm
+
+from coordinate import COORDINATES
+from coordinate import get_mask
+from utils.image_utils import save_image
+
+
+def parse_args():
+    """Parses arguments."""
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('jaco_path', type=str,
+                        help='Path to jacobian matrix.')
+    parser.add_argument('--region', type=str, default='eyes',
+                        help='The region to be used to compute jacobian.')
+    parser.add_argument('--save_dir', type=str, default='',
+                        help='Directory to save the results. If not specified,'
+                             'the results will be saved to '
+                            '`work_dirs/{TASK_SPECIFIC}/` by default')
+    parser.add_argument('--job', type=str, default='directions',
+                        help='Name for the job (default: directions)')
+    parser.add_argument('--name', type=str, default='resefa',
+                        help='Name of help save the results.')
+    parser.add_argument('--data_name', type=str, default='ffhq',
+                        help='Name of the dataset.')
+    parser.add_argument('--full_rank', action='store_true',
+                        help='Whether or not to full rank background'
+                             ' (default: False).')
+    parser.add_argument('--tao', type=float, default=1e-3,
+                        help='Coefficient to the identity matrix '
+                             '(default: 1e-3).')
+    return parser.parse_args()
+
+
+def main():
+    """Main function."""
+    args = parse_args()
+    assert os.path.exists(args.jaco_path)
+    Jacobians = np.load(args.jaco_path)
+    image_size = Jacobians.shape[2]
+    w_dim = Jacobians.shape[-1]
+    coord_dict = COORDINATES[args.data_name]
+    assert args.region in coord_dict, \
+        f'{args.region} coordinate is not defined in ' \
+        f'COORDINATE_{args.data_name}. Please define this region first!'
+    coords = coord_dict[args.region]
+    mask = get_mask(image_size, coordinate=coords)
+    foreground_ind = np.where(mask == 1)
+    background_ind = np.where((1 - mask) == 1)
+    temp_dir = f'./work_dirs/{args.job}/{args.data_name}/{args.region}'
+    save_dir = args.save_dir or temp_dir
+    os.makedirs(save_dir, exist_ok=True)
+    for ind in tqdm(range(Jacobians.shape[0])):
+        Jacobian = Jacobians[ind]
+        if len(Jacobian.shape) == 4:  # [H, W, 1, latent_dim]
+            Jaco_fore = Jacobian[foreground_ind[0], foreground_ind[1], 0]
+            Jaco_back = Jacobian[background_ind[0], background_ind[1], 0]
+        elif len(Jacobian.shape) == 5:  # [channel, H, W, 1, latent_dim]
+            Jaco_fore = Jacobian[:, foreground_ind[0], foreground_ind[1], 0]
+            Jaco_back = Jacobian[:, background_ind[0], background_ind[1], 0]
+        else:
+            raise ValueError('Shape of the Jacobian is not correct!')
+        Jaco_fore = np.reshape(Jaco_fore, [-1, w_dim])
+        Jaco_back = np.reshape(Jaco_back, [-1, w_dim])
+        coef_f = 1 / Jaco_fore.shape[0]
+        coef_b = 1 / Jaco_back.shape[0]
+        M_fore = coef_f * Jaco_fore.T.dot(Jaco_fore)
+        M_back = coef_b * Jaco_back.T.dot(Jaco_back)
+        if args.full_rank:
+            # J = J_b^TJ_b
+            # J = (J + tao * trace(J) * I)
+            print('Using full rank')
+            coef = args.tao * np.trace(M_back)
+            M_back = M_back + coef * np.identity(M_back.shape[0])
+        # inv(B) * A = lambda x
+        temp = np.linalg.inv(M_back).dot(M_fore)
+        eig_val, eig_vec = np.linalg.eig(temp)
+        eig_val = np.real(eig_val)
+        eig_vec = np.real(eig_vec)
+        directions = eig_vec.T
+        directions = directions[np.argsort(-eig_val)]
+        save_name = f'{save_dir}/image_{ind:02d}_region_{args.region}' \
+                    f'_name_{args.name}'
+        np.save(f'{save_name}.npy', directions)
+        mask_i = np.tile(mask[:, :, np.newaxis], [1, 1, 3]) * 255
+        save_image(f'{save_name}_mask.png', mask_i.astype(np.uint8))
+
+
+if __name__ == '__main__':
+    main()

compute_jacobian.py

@@ -0,0 +1,200 @@
+# python3.7
+"""Functions to compute Jacobian based on pre-trained GAN generator.
+
+Support StyleGAN2 or StyleGAN3
+"""
+
+import os
+import argparse
+import warnings
+from tqdm import tqdm
+import numpy as np
+
+import torch
+import torch.nn.functional as F
+from torch.autograd.functional import jacobian
+from models import build_model
+from utils.image_utils import save_image
+from utils.image_utils import postprocess_image
+from utils.custom_utils import to_numpy
+
+
+warnings.filterwarnings(action='ignore', category=UserWarning)
+
+
+def parse_args():
+    """Parses arguments."""
+    parser = argparse.ArgumentParser()
+    group = parser.add_argument_group('General options.')
+    group.add_argument('weight_path', type=str,
+                       help='Weight path to the pre-trained model.')
+    group.add_argument('--save_dir', type=str, default=None,
+                       help='Directory to save the results. If not specified, '
+                            'the results will be saved to '
+                            '`work_dirs/{TASK_SPECIFIC}/` by default.')
+    group.add_argument('--job', type=str, default='jacobians',
+                       help='Name for the job (default: jacobians)')
+    group.add_argument('--seed', type=int, default=4,
+                       help='Seed for sampling. (default: 4)')
+    group.add_argument('--nums', type=int, default=5,
+                       help='Number of samples to synthesized. (default: 5)')
+    group.add_argument('--img_size', type=int, default=1024,
+                       help='Size of the synthesized images. (default: 1024)')
+    group.add_argument('--w_dim', type=int, default=512,
+                       help='Dimension of the latent w. (default: 512)')
+    group.add_argument('--save_jpg', action='store_false',
+                       help='Whether to save the images used to compute '
+                            'jacobians. (default: True)')
+    group.add_argument('-d', '--data_name', type=str, default='ffhq',
+                       help='Name of the datasets. (default: ffhq)')
+    group.add_argument('--latent_path', type=str, default='',
+                       help='Path to the given latent codes. (default: None)')
+
+    group = parser.add_argument_group('StyleGAN2')
+    group.add_argument('--stylegan2', action='store_true',
+                       help='Whether or not using StyleGAN2. (default: False)')
+    group.add_argument('--scale_stylegan2', type=float, default=1.0,
+                       help='Scale for the number of channel fro stylegan2.')
+    group.add_argument('--randomize_noise', type=str, default='const',
+                       help='Noise type when computing. (const or random)')
+
+    group = parser.add_argument_group('StyleGAN3')
+    group.add_argument('--stylegan3', action='store_true',
+                       help='Whether or not using StyleGAN3. (default: False)')
+    group.add_argument('--cfg', type=str, default='T',
+                       help='Config of the stylegan3 (T/R).')
+    group.add_argument('--scale_stylegan3r', type=float, default=2.0,
+                       help='Scale for the number of channel for stylegan3 R.')
+    group.add_argument('--scale_stylegan3t', type=float, default=1.0,
+                       help='Scale for the number of channel for stylegan3 T.')
+    group.add_argument('--tx', type=float, default=0,
+                       help='Translate X-coordinate. (default: 0.0)')
+    group.add_argument('--ty', type=float, default=0,
+                       help='Translate Y-coordinate. (default: 0.0)')
+    group.add_argument('--rotate', type=float, default=0,
+                       help='Rotation angle in degrees. (default: 0)')
+
+    group = parser.add_argument_group('Jacobians')
+    group.add_argument('--b', type=float, default=1e-3,
+                       help='Constant when computing jacobians fast.')
+    group.add_argument('--batch_size', type=int, default=4,
+                       help='Batch size. (default: 4)')
+    return parser.parse_args()
+
+
+def main():
+    """Main function."""
+    args = parse_args()
+    # Parse model configuration.
+    assert (args.stylegan2 and not args.stylegan3) or \
+           (not args.stylegan2 and args.stylegan3)
+    job_disc = ''
+    if args.stylegan2:
+        config = dict(model_type='StyleGAN2Generator',
+                      resolution=args.img_size,
+                      w_dim=args.w_dim,
+                      fmaps_base=int(args.scale_stylegan2 * (32 << 10)),
+                      fmaps_max=512,)
+        job_disc += 'stylegan2'
+    else:
+        if args.stylegan3 and args.cfg == 'R':
+            config = dict(model_type='StyleGAN3Generator',
+                          resolution=args.img_size,
+                          w_dim=args.w_dim,
+                          fmaps_base=int(args.scale_stylegan3r * (32 << 10)),
+                          fmaps_max=1024,
+                          use_radial_filter=True,)
+            job_disc += 'stylegan3r'
+        elif args.stylegan3 and args.cfg == 'T':
+            config = dict(model_type='StyleGAN3Generator',
+                          resolution=args.img_size,
+                          w_dim=args.w_dim,
+                          fmaps_base=int(args.scale_stylegan3t * (32 << 10)),
+                          fmaps_max=512,
+                          use_radial_filter=False,
+                          kernel_size=3,)
+            job_disc += 'stylegan3t'
+        else:
+            raise TypeError(f'StyleGAN3 config type error, need `R/T`,'
+                            f' but got {args.cfg}')
+    job_name = f'seed_{args.seed}_num_{args.nums}_{job_disc}'
+    temp_dir = f'work_dirs/{args.job}/{args.data_name}/{job_name}'
+    save_dir = args.save_dir or temp_dir
+    os.makedirs(save_dir, exist_ok=True)
+    if args.save_jpg:
+        os.makedirs(f'{save_dir}/images', exist_ok=True)
+
+    print('Building generator...')
+    generator = build_model(**config)
+    checkpoint_path = args.weight_path
+    print(f'Loading checkpoint from `{checkpoint_path}` ...')
+    checkpoint = torch.load(checkpoint_path, map_location='cpu')['models']
+    if 'generator_smooth' in checkpoint:
+        generator.load_state_dict(checkpoint['generator_smooth'])
+    else:
+        generator.load_state_dict(checkpoint['generator'])
+    generator = generator.eval().cuda()
+    print('Finish loading checkpoint.')
+
+    # Set random seed.
+    np.random.seed(args.seed)
+    torch.manual_seed(args.seed)
+    if os.path.exists(args.latent_path):
+        latent_zs = np.load(args.latent_path)
+        latent_zs = latent_zs[:args.nums]
+    else:
+        latent_zs = np.random.randn(args.nums, generator.z_dim)
+    latent_zs = torch.from_numpy(latent_zs.astype(np.float32))
+    latent_zs = latent_zs.cuda()
+    with torch.no_grad():
+        latent_ws = generator.mapping(latent_zs)['w']
+    print(f'Shape of the latent w: {latent_ws.shape}')
+
+    def syn2jaco(w):
+        """Wrap the synthesized function to compute the Jacobian easily.
+
+        Basically, this function defines a generator that takes the input
+        from the W space and then synthesizes an image. If the image is
+        larger than 256, it will be resized to 256 to save the time and
+        storage.
+
+        Args:
+            w: latent code from the W space
+
+        Returns:
+            An image with the size of [1, 256, 256]
+        """
+        wp = w.unsqueeze(1).repeat((1, generator.num_layers, 1))
+        image = generator.synthesis(wp)['image']
+        if image.shape[-1] > 256:
+            scale = 256 / image.shape[-1]
+            image = F.interpolate(image, scale_factor=scale)
+            image = torch.sum(image, dim=1)
+        return image
+
+    jacobians = []
+    for idx in tqdm(range(latent_zs.shape[0])):
+        latent_w = latent_ws[idx:idx+1]
+        jac_i = jacobian(func=syn2jaco,
+                            inputs=latent_w,
+                            create_graph=False,
+                            strict=False)
+        jacobians.append(jac_i)
+        if args.save_jpg:
+            wp = latent_w.unsqueeze(1).repeat((1, generator.num_layers, 1))
+            syn_outputs = generator.synthesis(wp)['image']
+            syn_outputs = to_numpy(syn_outputs)
+            images = postprocess_image(syn_outputs)
+            save_path = f'{save_dir}/images/{idx:06d}.jpg'
+            save_image(save_path, images[0])
+    jacobians = torch.cat(jacobians, dim=0)
+    jacobians = to_numpy(jacobians)
+    print(f'shape of the jacobian: {jacobians.shape}')
+    latent_ws = to_numpy(latent_ws)
+    np.save(f'{save_dir}/latent_codes.npy', latent_ws)
+    np.save(f'{save_dir}/jacobians_w.npy', jacobians)
+    print(f'Finish computing {args.nums} jacobians.')
+
+
+if __name__ == '__main__':
+    main()

coordinate.py

@@ -0,0 +1,142 @@
+# python3.7
+"""Utility functions to help define the region coordinates within an image."""
+
+import os
+from glob import glob
+import argparse
+import numpy as np
+import cv2
+from tqdm import tqdm
+from utils.parsing_utils import parse_index
+
+
+def get_mask_by_coordinates(image_size, coordinate):
+    """Get mask using the provided coordinates."""
+    mask = np.zeros([image_size, image_size], dtype=np.float32)
+    center_x, center_y = coordinate[0], coordinate[1]
+    crop_x, crop_y = coordinate[2], coordinate[3]
+    xx = center_x - crop_x // 2
+    yy = center_y - crop_y // 2
+    mask[xx:xx + crop_x, yy:yy + crop_y] = 1.
+    return mask
+
+
+def get_mask_by_segmentation(seg_mask, label):
+    """Get the mask using the segmentation array and labels."""
+    zeros = np.zeros_like(seg_mask)
+    ones = np.ones_like(seg_mask)
+    mask = np.where(seg_mask == label, ones, zeros)
+    return mask
+
+
+def get_mask(image_size, coordinate=None, seg_mask=None, labels='1'):
+    """Get mask using either the coordinate or the segmentation array."""
+    if coordinate is not None:
+        print('Using coordinate to get mask!')
+        mask = get_mask_by_coordinates(image_size, coordinate)
+    else:
+        print('Using segmentation to get the mask!')
+        print(f'Using label {labels}')
+        mask = np.zeros_like(seg_mask)
+        for label_ in labels:
+            mask += get_mask_by_segmentation(seg_mask, int(label_))
+        mask = np.clip(mask, a_min=0, a_max=1)
+
+    return mask
+
+
+# For FFHQ [center_x, center_y, height, width]
+# Those coordinates are suitable for both ffhq and metface.
+COORDINATE_ffhq = {'left_eye': [120, 95, 20, 38],
+                   'right_eye': [120, 159, 20, 38],
+                   'eyes': [120, 128, 20, 115],
+                   'nose': [142, 131, 40, 46],
+                   'mouth': [184, 127, 30, 70],
+                   'chin': [217, 130, 42, 110],
+                   'eyebrow': [126, 105, 15, 118],
+                   }
+
+
+# For FFHQ unaligned
+COORDINATE_ffhqu = {'eyesr2': [134, 116, 30, 115],
+                    'eyesr3': [64, 128, 26, 115],
+                    'eyest0': [70, 88, 30, 115],
+                    'eyest3': [108, 142, 26, 115],
+                    }
+
+# [center_x, center_y, height, width]
+COORDINATE_biggan = {'center0': [120, 120, 80, 80],
+                     'center1': [120, 120, 130, 130],
+                     'center2': [120, 120,  200, 200],
+                     'left_side': [128, 64, 256, 128],
+                     'top_side': [64, 128, 128, 256],
+                     'head0': [89, 115, 49, 70],
+                     'head1': [93, 110, 48, 70]}
+
+
+COORDINATES = {'ffhq': COORDINATE_ffhq,
+               'ffhqu': COORDINATE_ffhqu,
+               'biggan': COORDINATE_biggan
+               }
+
+
+def parse_args():
+    """Parses arguments."""
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--image_path', type=str, default='',
+                        help='The path to the image.')
+    parser.add_argument('--mask_path', type=str, default='',
+                        help='The path to the mask.')
+    parser.add_argument('--save_dir', type=str, default='',
+                        help='The path to the image.')
+    parser.add_argument('--label', type=str, default=None,
+                        help='The label number in the mask.')
+    parser.add_argument('--data', type=str, default='ffhq',
+                        help='The name of the dataset to test.')
+    parser.add_argument('--num', type=int, default=0,
+                        help='number of image to display.')
+    parser.add_argument('--img_type', type=str, default='jpeg',
+                        help='Format of the image.')
+
+    return parser.parse_args()
+
+
+def main():
+    """Main function to show an image with masks"""
+    args = parse_args()
+    save_dir = args.save_dir or './temp_mask'
+    os.makedirs(save_dir, exist_ok=True)
+    images = sorted(glob(f'{args.image_path}/*.{args.img_type}'))[args.num:]
+    label_files = sorted(glob(f'{args.mask_path}/*.npy'))[args.num:]
+    COORDINATE = COORDINATES[args.data]
+    for i, image in tqdm(enumerate(images)):
+        img = cv2.imread(image)
+        im_name = image.split('/')[-1].split('.')[0]
+        if args.label is None:
+            for name, coord in COORDINATE.items():
+                if len(coord) == 0:
+                    continue
+                mask = np.zeros(img.shape, dtype=np.float32)
+                center_x, center_y = coord[0], coord[1]
+                crop_x, crop_y = coord[2], coord[3]
+                xx = center_x - crop_x // 2
+                yy = center_y - crop_y // 2
+                mask[xx:xx + crop_x, yy:yy + crop_y, :] = 1.
+                img_ = img * mask
+                cv2.imwrite(f'{save_dir}/{im_name}_{name}.png', img_)
+        else:
+            print('Using segmentation to get the mask!')
+            seg_mask = np.load(label_files[i])
+            labels = parse_index(args.label)
+            print(f'Using label {labels}')
+            mask = np.zeros_like(seg_mask)
+            for label_ in labels:
+                mask += get_mask_by_segmentation(seg_mask, int(label_))
+            mask = np.clip(mask, a_min=0, a_max=1)
+            img_ = img * mask[:, :, np.newaxis]
+            cv2.imwrite(f'{save_dir}/{im_name}_{args.label}.png', img_)
+
+
+if __name__ == '__main__':
+    main()

directions/afhq/stylegan3/eyes-r.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:959982185cd0401b8ab984ad2c1c22c01a494d8928a765e353c4fc34e9b079a2
+size 2176

directions/ffhq/stylegan2/eyebrows.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f8a7385ee951fa93b34044a768254a937327c87b0d86a569b62ae4593ae0b765
+size 2176

directions/ffhq/stylegan2/eyesize.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:74d417489944e2f2c6ee9622c4a3745edf20edad5837f5d4b2d89847003d0ec5
+size 2176

directions/ffhq/stylegan2/gaze_direction.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:98c7787363716096706db4f6cc039cf5ffc773142a0f790ed185300ddaf8fc0e
+size 2176

directions/ffhq/stylegan2/lipstick.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9ce301bb70f35e10c06b9d6fbb485644a65ebf1ac57a080e70fc3448b3343dc6
+size 2176

directions/ffhq/stylegan2/mouth.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5cf4fe4565fd2d3f363d0b2e7de1073f1379c8488113b99f8d3bad9ba53a0fa4
+size 2176

directions/ffhq/stylegan2/nose_length.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7a7d3685b38f18e87a49cc3a79106f18dbdf30b1e8d3db5df21c84d592566ea5
+size 2176

directions/ffhq/stylegan3/eyes-r.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:91f7c4432a030568b1aac414c30e7f00cf4d441228a625eb8d971f5f28b036db
+size 2176

manipulate.py

@@ -0,0 +1,253 @@
+# python3.7
+"""Manipulates synthesized or real images with existing boundary.
+
+Support StyleGAN2 and StyleGAN3.
+"""
+
+import os.path
+import argparse
+import numpy as np
+from tqdm import tqdm
+import torch
+
+from models import build_model
+from utils.visualizers.html_visualizer import HtmlVisualizer
+from utils.image_utils import save_image
+from utils.parsing_utils import parse_index
+from utils.image_utils import postprocess_image
+from utils.custom_utils import to_numpy, linear_interpolate
+from utils.custom_utils import make_transform
+
+
+def parse_args():
+    """Parses arguments."""
+    parser = argparse.ArgumentParser()
+    group = parser.add_argument_group('General options.')
+    group.add_argument('weight_path', type=str,
+                       help='Weight path to the pre-trained model.')
+    group.add_argument('boundary_path', type=str,
+                       help='Path to the attribute vectors.')
+    group.add_argument('--save_dir', type=str, default=None,
+                       help='Directory to save the results. If not specified, '
+                            'the results will be saved to '
+                            '`work_dirs/{TASK_SPECIFIC}/` by default.')
+    group.add_argument('--job', type=str, default='manipulations',
+                       help='Name for the job. (default: manipulations)')
+    group.add_argument('--seed', type=int, default=4,
+                       help='Seed for sampling. (default: 4)')
+    group.add_argument('--nums', type=int, default=10,
+                       help='Number of samples to synthesized. (default: 10)')
+    group.add_argument('--img_size', type=int, default=1024,
+                       help='Size of the synthesized images. (default: 1024)')
+    group.add_argument('--vis_size', type=int, default=256,
+                       help='Size of the visualize images. (default: 256)')
+    group.add_argument('--w_dim', type=int, default=512,
+                       help='Dimension of the latent w. (default: 512)')
+    group.add_argument('--batch_size', type=int, default=4,
+                       help='Batch size. (default: 4)')
+    group.add_argument('--save_jpg', action='store_true', default=False,
+                       help='Whether to save raw image. (default: False)')
+    group.add_argument('-d', '--data_name', type=str, default='ffhq',
+                       help='Name of the datasets. (default: ffhq)')
+    group.add_argument('--latent_path', type=str, default='',
+                       help='Path to the given latent codes. (default: None)')
+    group.add_argument('--trunc_psi', type=float, default=0.7,
+                       help='Psi factor used for truncation. (default: 0.7)')
+    group.add_argument('--trunc_layers', type=int, default=8,
+                       help='Number of layers to perform truncation.'
+                            ' (default: 8)')
+    group.add_argument('--name', type=str, default='resefa',
+                       help='Name of help save the results.')
+
+    group = parser.add_argument_group('StyleGAN2')
+    group.add_argument('--stylegan2', action='store_true',
+                       help='Whether or not using StyleGAN2. (default: False)')
+    group.add_argument('--scale_stylegan2', type=float, default=1.0,
+                       help='Scale for the number of channel fro stylegan2.')
+    group.add_argument('--randomize_noise', type=str, default='const',
+                       help='Noise type when editing. (const or random)')
+
+    group = parser.add_argument_group('StyleGAN3')
+    group.add_argument('--stylegan3', action='store_true',
+                       help='Whether or not using StyleGAN3. (default: False)')
+    group.add_argument('--cfg', type=str, default='T',
+                       help='Config of the stylegan3 (T/R)')
+    group.add_argument('--scale_stylegan3r', type=float, default=2.0,
+                       help='Scale for the number of channel for stylegan3 R.')
+    group.add_argument('--scale_stylegan3t', type=float, default=1.0,
+                       help='Scale for the number of channel for stylegan3 T.')
+    group.add_argument('--tx', type=float, default=0,
+                       help='Translate X-coordinate. (default: 0.0)')
+    group.add_argument('--ty', type=float, default=0,
+                       help='Translate Y-coordinate. (default: 0.0)')
+    group.add_argument('--rotate', type=float, default=0,
+                       help='Rotation angle in degrees. (default: 0)')
+
+    group = parser.add_argument_group('Manipulation')
+    group.add_argument('--mani_layers', type=str, default='4,5,6,7',
+                       help='The layers will be manipulated.'
+                            '(default: 4,5,6,7). For the eyebrow and lipstick,'
+                            'using [8-11] layers instead.')
+    group.add_argument('--step', type=int, default=7,
+                       help='Number of manipulation steps. (default: 7)')
+    group.add_argument('--start', type=int, default=0,
+                       help='The start index of the manipulation directions.')
+    group.add_argument('--end', type=int, default=1,
+                       help='The end index of the manipulation directions.')
+    group.add_argument('--start_distance', type=float, default=-10.0,
+                       help='Start distance for manipulation. (default: -10.0)')
+    group.add_argument('--end_distance', type=float, default=10.0,
+                       help='End distance for manipulation. (default: 10.0)')
+
+    return parser.parse_args()
+
+
+def main():
+    """Main function."""
+    args = parse_args()
+    # Parse model configuration.
+    assert (args.stylegan2 and not args.stylegan3) or \
+           (not args.stylegan2 and args.stylegan3)
+    checkpoint_path = args.weight_path
+    boundary_path = args.boundary_path
+    assert os.path.exists(checkpoint_path)
+    assert os.path.exists(boundary_path)
+    boundary_name = os.path.splitext(os.path.basename(boundary_path))[0]
+    job_disc = ''
+    if args.stylegan2:
+        config = dict(model_type='StyleGAN2Generator',
+                      resolution=args.img_size,
+                      w_dim=args.w_dim,
+                      fmaps_base=int(args.scale_stylegan2 * (32 << 10)),
+                      fmaps_max=512,)
+        job_disc += 'stylegan2'
+    else:
+        if args.stylegan3 and args.cfg == 'R':
+            config = dict(model_type='StyleGAN3Generator',
+                          resolution=args.img_size,
+                          w_dim=args.w_dim,
+                          fmaps_base=int(args.scale_stylegan3r * (32 << 10)),
+                          fmaps_max=1024,
+                          use_radial_filter=True,)
+            job_disc += 'stylegan3r'
+        elif args.stylegan3 and args.cfg == 'T':
+            config = dict(model_type='StyleGAN3Generator',
+                          resolution=args.img_size,
+                          w_dim=args.w_dim,
+                          fmaps_base=int(args.scale_stylegan3t * (32 << 10)),
+                          fmaps_max=512,
+                          use_radial_filter=False,
+                          kernel_size=3,)
+            job_disc += 'stylegan3t'
+        else:
+            raise TypeError(f'StyleGAN3 config type error, need `R/T`,'
+                            f' but got {args.cfg} instead.')
+
+    # Get work directory and job name.
+    save_dir = args.save_dir or f'work_dirs/{args.job}/{args.data_name}'
+    os.makedirs(save_dir, exist_ok=True)
+    job_name = f'seed_{args.seed}_num_{args.nums}_{job_disc}_{boundary_name}'
+    os.makedirs(f'{save_dir}/{job_name}', exist_ok=True)
+
+    print('Building generator...')
+    generator = build_model(**config)
+    print(f'Loading checkpoint from `{checkpoint_path}` ...')
+    checkpoint = torch.load(checkpoint_path, map_location='cpu')['models']
+    if 'generator_smooth' in checkpoint:
+        generator.load_state_dict(checkpoint['generator_smooth'])
+    else:
+        generator.load_state_dict(checkpoint['generator'])
+    generator = generator.eval().cuda()
+    print('Finish loading checkpoint.')
+    if args.stylegan3 and hasattr(generator.synthesis, 'early_layer'):
+        m = make_transform(args.tx, args.ty, args.rotate)
+        m = np.linalg.inv(m)
+        generator.synthesis.early_layer.transform.copy_(torch.from_numpy(m))
+
+    np.random.seed(args.seed)
+    torch.manual_seed(args.seed)
+    if os.path.exists(args.latent_path):
+        print(f'Load latent codes from {args.latent_path}')
+        latent_zs = np.load(args.latent_path)
+        latent_zs = latent_zs[:args.nums]
+    else:
+        print('Sampling latent code randomly')
+        latent_zs = np.random.randn(args.nums, generator.z_dim)
+    latent_zs = torch.from_numpy(latent_zs.astype(np.float32))
+    latent_zs = latent_zs.cuda()
+    num_images = latent_zs.shape[0]
+    wp = []
+    for idx in range(0, num_images, args.batch_size):
+        latent_z = latent_zs[idx:idx+args.batch_size]
+        latent_w_ = generator.mapping(latent_z, None)['wp']
+        wp.append(latent_w_)
+    wp = torch.cat(wp, dim=0)
+    trunc_psi = args.trunc_psi
+    trunc_layers = args.trunc_layers
+    if trunc_psi < 1.0 and trunc_layers > 0:
+        w_avg = generator.w_avg
+        w_avg = w_avg.reshape(1, -1, generator.w_dim)[:, :trunc_layers]
+        wp[:, :trunc_layers] = w_avg.lerp(wp[:, :trunc_layers], trunc_psi)
+    print(f'Shape of the latent ws: {wp.shape}')
+    image_list = []
+    for i in range(num_images):
+        image_list.append(f'{i:06d}')
+
+    print('Loading boundary.')
+    directions = np.load(boundary_path)
+    layer_index = parse_index(args.mani_layers)
+    if not layer_index:
+        layer_index = list(range(generator.num_layers - 1))
+    print(f'Manipulating on layers `{layer_index}`.')
+
+    vis_size = None if args.vis_size == 0 else args.vis_size
+    delta_num = args.end - args.start
+    visualizer = HtmlVisualizer(num_rows=num_images * delta_num,
+                                num_cols=args.step + 2,
+                                image_size=vis_size)
+    visualizer.set_headers(
+        ['Name', 'Origin'] +
+        [f'Step {i:02d}' for i in range(1, args.step + 1)]
+    )
+    # Manipulate images.
+    print('Start manipulation.')
+    for row in tqdm(range(num_images)):
+        latent_w = wp[row:row+1]
+        images_ori = generator.synthesis(latent_w)['image']
+        images_ori = postprocess_image(to_numpy(images_ori))
+        if args.save_jpg:
+            save_image(f'{save_dir}/{job_name}/{row:06d}_orin.jpg',
+                       images_ori[0])
+        for num_direc in range(args.start, args.end):
+            html_row = num_direc - args.start
+            direction = directions[num_direc:num_direc+1]
+            direction = np.tile(direction, [1, generator.num_layers, 1])
+            visualizer.set_cell(row * delta_num + html_row, 0,
+                                text=f'{image_list[row]}_{num_direc:03d}')
+            visualizer.set_cell(row * delta_num + html_row, 1,
+                                image=images_ori[0])
+            mani_codes = linear_interpolate(latent_code=to_numpy(latent_w),
+                                            boundary=direction,
+                                            layer_index=layer_index,
+                                            start_distance=args.start_distance,
+                                            end_distance=args.end_distance,
+                                            steps=args.step)
+            mani_codes = torch.from_numpy(mani_codes.astype(np.float32)).cuda()
+            for idx in range(0, mani_codes.shape[0], args.batch_size):
+                codes_ = mani_codes[idx:idx+args.batch_size]
+                images_ = generator.synthesis(codes_)['image']
+                images_ = postprocess_image(to_numpy(images_))
+                for i in range(images_.shape[0]):
+                    visualizer.set_cell(row * delta_num + html_row, idx+i+2,
+                                        image=images_[i])
+                    if args.save_jpg:
+                        save_image(f'{save_dir}/{job_name}/{row:06d}_ind_'
+                                   f'{num_direc:06d}_mani_{idx+i:06d}.jpg',
+                                   images_[i])
+    # Save results.
+    np.save(f'{save_dir}/{job_name}/latent_codes.npy', to_numpy(wp))
+    visualizer.save(f'{save_dir}/{job_name}_{args.name}.html')
+
+
+if __name__ == '__main__':
+    main()

models/__init__.py

@@ -0,0 +1,45 @@
+# python3.7
+"""Collects all models."""
+
+from .pggan_generator import PGGANGenerator
+from .pggan_discriminator import PGGANDiscriminator
+from .stylegan_generator import StyleGANGenerator
+from .stylegan_discriminator import StyleGANDiscriminator
+from .stylegan2_generator import StyleGAN2Generator
+from .stylegan2_discriminator import StyleGAN2Discriminator
+from .stylegan3_generator import StyleGAN3Generator
+from .ghfeat_encoder import GHFeatEncoder
+from .perceptual_model import PerceptualModel
+from .inception_model import InceptionModel
+
+__all__ = ['build_model']
+
+_MODELS = {
+    'PGGANGenerator': PGGANGenerator,
+    'PGGANDiscriminator': PGGANDiscriminator,
+    'StyleGANGenerator': StyleGANGenerator,
+    'StyleGANDiscriminator': StyleGANDiscriminator,
+    'StyleGAN2Generator': StyleGAN2Generator,
+    'StyleGAN2Discriminator': StyleGAN2Discriminator,
+    'StyleGAN3Generator': StyleGAN3Generator,
+    'GHFeatEncoder': GHFeatEncoder,
+    'PerceptualModel': PerceptualModel.build_model,
+    'InceptionModel': InceptionModel.build_model
+}
+
+
+def build_model(model_type, **kwargs):
+    """Builds a model based on its class type.
+
+    Args:
+        model_type: Class type to which the model belongs, which is case
+            sensitive.
+        **kwargs: Additional arguments to build the model.
+
+    Raises:
+        ValueError: If the `model_type` is not supported.
+    """
+    if model_type not in _MODELS:
+        raise ValueError(f'Invalid model type: `{model_type}`!\n'
+                         f'Types allowed: {list(_MODELS)}.')
+    return _MODELS[model_type](**kwargs)

models/ghfeat_encoder.py

@@ -0,0 +1,563 @@
+# python3.7
+"""Contains the implementation of encoder used in GH-Feat (including IDInvert).
+
+ResNet is used as the backbone.
+
+GH-Feat paper: https://arxiv.org/pdf/2007.10379.pdf
+IDInvert paper: https://arxiv.org/pdf/2004.00049.pdf
+
+NOTE: Please use `latent_num` and `num_latents_per_head` to control the
+inversion space, such as Y-space used in GH-Feat and W-space used in IDInvert.
+In addition, IDInvert sets `use_fpn` and `use_sam` as `False` by default.
+"""
+
+import numpy as np
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.distributed as dist
+
+__all__ = ['GHFeatEncoder']
+
+# Resolutions allowed.
+_RESOLUTIONS_ALLOWED = [8, 16, 32, 64, 128, 256, 512, 1024]
+
+# pylint: disable=missing-function-docstring
+
+class BasicBlock(nn.Module):
+    """Implementation of ResNet BasicBlock."""
+
+    expansion = 1
+
+    def __init__(self,
+                 inplanes,
+                 planes,
+                 base_width=64,
+                 stride=1,
+                 groups=1,
+                 dilation=1,
+                 norm_layer=None,
+                 downsample=None):
+        super().__init__()
+        if base_width != 64:
+            raise ValueError(f'BasicBlock of ResNet only supports '
+                             f'`base_width=64`, but {base_width} received!')
+        if stride not in [1, 2]:
+            raise ValueError(f'BasicBlock of ResNet only supports `stride=1` '
+                             f'and `stride=2`, but {stride} received!')
+        if groups != 1:
+            raise ValueError(f'BasicBlock of ResNet only supports `groups=1`, '
+                             f'but {groups} received!')
+        if dilation != 1:
+            raise ValueError(f'BasicBlock of ResNet only supports '
+                             f'`dilation=1`, but {dilation} received!')
+        assert self.expansion == 1
+
+        self.stride = stride
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        self.conv1 = nn.Conv2d(in_channels=inplanes,
+                               out_channels=planes,
+                               kernel_size=3,
+                               stride=stride,
+                               padding=1,
+                               groups=1,
+                               dilation=1,
+                               bias=False)
+        self.bn1 = norm_layer(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = nn.Conv2d(in_channels=planes,
+                               out_channels=planes,
+                               kernel_size=3,
+                               stride=1,
+                               padding=1,
+                               groups=1,
+                               dilation=1,
+                               bias=False)
+        self.bn2 = norm_layer(planes)
+        self.downsample = downsample
+
+    def forward(self, x):
+        identity = self.downsample(x) if self.downsample is not None else x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out + identity)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    """Implementation of ResNet Bottleneck."""
+
+    expansion = 4
+
+    def __init__(self,
+                 inplanes,
+                 planes,
+                 base_width=64,
+                 stride=1,
+                 groups=1,
+                 dilation=1,
+                 norm_layer=None,
+                 downsample=None):
+        super().__init__()
+        if stride not in [1, 2]:
+            raise ValueError(f'Bottleneck of ResNet only supports `stride=1` '
+                             f'and `stride=2`, but {stride} received!')
+
+        width = int(planes * (base_width / 64)) * groups
+        self.stride = stride
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        self.conv1 = nn.Conv2d(in_channels=inplanes,
+                               out_channels=width,
+                               kernel_size=1,
+                               stride=1,
+                               padding=0,
+                               dilation=1,
+                               groups=1,
+                               bias=False)
+        self.bn1 = norm_layer(width)
+        self.conv2 = nn.Conv2d(in_channels=width,
+                               out_channels=width,
+                               kernel_size=3,
+                               stride=stride,
+                               padding=dilation,
+                               groups=groups,
+                               dilation=dilation,
+                               bias=False)
+        self.bn2 = norm_layer(width)
+        self.conv3 = nn.Conv2d(in_channels=width,
+                               out_channels=planes * self.expansion,
+                               kernel_size=1,
+                               stride=1,
+                               padding=0,
+                               dilation=1,
+                               groups=1,
+                               bias=False)
+        self.bn3 = norm_layer(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+
+    def forward(self, x):
+        identity = self.downsample(x) if self.downsample is not None else x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+        out = self.relu(out + identity)
+
+        return out
+
+
+class GHFeatEncoder(nn.Module):
+    """Define the ResNet-based encoder network for GAN inversion.
+
+    On top of the backbone, there are several task-heads to produce inverted
+    codes. Please use `latent_dim` and `num_latents_per_head` to define the
+    structure. For example, `latent_dim = [512] * 14` and
+    `num_latents_per_head = [4, 4, 6]` can be used for StyleGAN inversion with
+    14-layer latent codes, where 3 task heads (corresponding to 4, 4, 6 layers,
+    respectively) are used.
+
+    Settings for the encoder network:
+
+    (1) resolution: The resolution of the output image.
+    (2) latent_dim: Dimension of the latent space. A number (one code will be
+        produced), or a list of numbers regarding layer-wise latent codes.
+    (3) num_latents_per_head: Number of latents that is produced by each head.
+    (4) image_channels: Number of channels of the output image. (default: 3)
+    (5) final_res: Final resolution of the convolutional layers. (default: 4)
+
+    ResNet-related settings:
+
+    (1) network_depth: Depth of the network, like 18 for ResNet18. (default: 18)
+    (2) inplanes: Number of channels of the first convolutional layer.
+        (default: 64)
+    (3) groups: Groups of the convolution, used in ResNet. (default: 1)
+    (4) width_per_group: Number of channels per group, used in ResNet.
+        (default: 64)
+    (5) replace_stride_with_dilation: Whether to replace stride with dilation,
+        used in ResNet. (default: None)
+    (6) norm_layer: Normalization layer used in the encoder. If set as `None`,
+        `nn.BatchNorm2d` will be used. Also, please NOTE that when using batch
+        normalization, the batch size is required to be larger than one for
+        training. (default: nn.BatchNorm2d)
+    (7) max_channels: Maximum number of channels in each layer. (default: 512)
+
+    Task-head related settings:
+
+    (1) use_fpn: Whether to use Feature Pyramid Network (FPN) before outputting
+        the latent code. (default: True)
+    (2) fpn_channels: Number of channels used in FPN. (default: 512)
+    (3) use_sam: Whether to use Spatial Alignment Module (SAM) before outputting
+        the latent code. (default: True)
+    (4) sam_channels: Number of channels used in SAM. (default: 512)
+    """
+
+    arch_settings = {
+        18: (BasicBlock,  [2, 2, 2, 2]),
+        34: (BasicBlock,  [3, 4, 6, 3]),
+        50: (Bottleneck,  [3, 4, 6, 3]),
+        101: (Bottleneck, [3, 4, 23, 3]),
+        152: (Bottleneck, [3, 8, 36, 3])
+    }
+
+    def __init__(self,
+                 resolution,
+                 latent_dim,
+                 num_latents_per_head,
+                 image_channels=3,
+                 final_res=4,
+                 network_depth=18,
+                 inplanes=64,
+                 groups=1,
+                 width_per_group=64,
+                 replace_stride_with_dilation=None,
+                 norm_layer=nn.BatchNorm2d,
+                 max_channels=512,
+                 use_fpn=True,
+                 fpn_channels=512,
+                 use_sam=True,
+                 sam_channels=512):
+        super().__init__()
+
+        if resolution not in _RESOLUTIONS_ALLOWED:
+            raise ValueError(f'Invalid resolution: `{resolution}`!\n'
+                             f'Resolutions allowed: {_RESOLUTIONS_ALLOWED}.')
+        if network_depth not in self.arch_settings:
+            raise ValueError(f'Invalid network depth: `{network_depth}`!\n'
+                             f'Options allowed: '
+                             f'{list(self.arch_settings.keys())}.')
+        if isinstance(latent_dim, int):
+            latent_dim = [latent_dim]
+        assert isinstance(latent_dim, (list, tuple))
+        assert isinstance(num_latents_per_head, (list, tuple))
+        assert sum(num_latents_per_head) == len(latent_dim)
+
+        self.resolution = resolution
+        self.latent_dim = latent_dim
+        self.num_latents_per_head = num_latents_per_head
+        self.num_heads = len(self.num_latents_per_head)
+        self.image_channels = image_channels
+        self.final_res = final_res
+        self.inplanes = inplanes
+        self.network_depth = network_depth
+        self.groups = groups
+        self.dilation = 1
+        self.base_width = width_per_group
+        self.replace_stride_with_dilation = replace_stride_with_dilation
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        if norm_layer == nn.BatchNorm2d and dist.is_initialized():
+            norm_layer = nn.SyncBatchNorm
+        self.norm_layer = norm_layer
+        self.max_channels = max_channels
+        self.use_fpn = use_fpn
+        self.fpn_channels = fpn_channels
+        self.use_sam = use_sam
+        self.sam_channels = sam_channels
+
+        block_fn, num_blocks_per_stage = self.arch_settings[network_depth]
+
+        self.num_stages = int(np.log2(resolution // final_res)) - 1
+        # Add one block for additional stages.
+        for i in range(len(num_blocks_per_stage), self.num_stages):
+            num_blocks_per_stage.append(1)
+        if replace_stride_with_dilation is None:
+            replace_stride_with_dilation = [False] * self.num_stages
+
+        # Backbone.
+        self.conv1 = nn.Conv2d(in_channels=self.image_channels,
+                               out_channels=self.inplanes,
+                               kernel_size=7,
+                               stride=2,
+                               padding=3,
+                               bias=False)
+        self.bn1 = norm_layer(self.inplanes)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+
+        self.stage_channels = [self.inplanes]
+        self.stages = nn.ModuleList()
+        for i in range(self.num_stages):
+            inplanes = self.inplanes if i == 0 else planes * block_fn.expansion
+            planes = min(self.max_channels, self.inplanes * (2 ** i))
+            num_blocks = num_blocks_per_stage[i]
+            stride = 1 if i == 0 else 2
+            dilate = replace_stride_with_dilation[i]
+            self.stages.append(self._make_stage(block_fn=block_fn,
+                                                inplanes=inplanes,
+                                                planes=planes,
+                                                num_blocks=num_blocks,
+                                                stride=stride,
+                                                dilate=dilate))
+            self.stage_channels.append(planes * block_fn.expansion)
+
+        if self.num_heads > len(self.stage_channels):
+            raise ValueError('Number of task heads is larger than number of '
+                             'stages! Please reduce the number of heads.')
+
+        # Task-head.
+        if self.num_heads == 1:
+            self.use_fpn = False
+            self.use_sam = False
+
+        if self.use_fpn:
+            fpn_pyramid_channels = self.stage_channels[-self.num_heads:]
+            self.fpn = FPN(pyramid_channels=fpn_pyramid_channels,
+                           out_channels=self.fpn_channels)
+        if self.use_sam:
+            if self.use_fpn:
+                sam_pyramid_channels = [self.fpn_channels] * self.num_heads
+            else:
+                sam_pyramid_channels = self.stage_channels[-self.num_heads:]
+            self.sam = SAM(pyramid_channels=sam_pyramid_channels,
+                           out_channels=self.sam_channels)
+
+        self.heads = nn.ModuleList()
+        for head_idx in range(self.num_heads):
+            # Parse in_channels.
+            if self.use_sam:
+                in_channels = self.sam_channels
+            elif self.use_fpn:
+                in_channels = self.fpn_channels
+            else:
+                in_channels = self.stage_channels[head_idx - self.num_heads]
+            in_channels = in_channels * final_res * final_res
+
+            # Parse out_channels.
+            start_latent_idx = sum(self.num_latents_per_head[:head_idx])
+            end_latent_idx = sum(self.num_latents_per_head[:head_idx + 1])
+            out_channels = sum(self.latent_dim[start_latent_idx:end_latent_idx])
+
+            self.heads.append(CodeHead(in_channels=in_channels,
+                                       out_channels=out_channels,
+                                       norm_layer=self.norm_layer))
+
+    def _make_stage(self,
+                    block_fn,
+                    inplanes,
+                    planes,
+                    num_blocks,
+                    stride,
+                    dilate):
+        norm_layer = self.norm_layer
+        downsample = None
+        previous_dilation = self.dilation
+        if dilate:
+            self.dilation *= stride
+            stride = 1
+        if stride != 1 or inplanes != planes * block_fn.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(in_channels=inplanes,
+                          out_channels=planes * block_fn.expansion,
+                          kernel_size=1,
+                          stride=stride,
+                          padding=0,
+                          dilation=1,
+                          groups=1,
+                          bias=False),
+                norm_layer(planes * block_fn.expansion),
+            )
+
+        blocks = []
+        blocks.append(block_fn(inplanes=inplanes,
+                               planes=planes,
+                               base_width=self.base_width,
+                               stride=stride,
+                               groups=self.groups,
+                               dilation=previous_dilation,
+                               norm_layer=norm_layer,
+                               downsample=downsample))
+        for _ in range(1, num_blocks):
+            blocks.append(block_fn(inplanes=planes * block_fn.expansion,
+                                   planes=planes,
+                                   base_width=self.base_width,
+                                   stride=1,
+                                   groups=self.groups,
+                                   dilation=self.dilation,
+                                   norm_layer=norm_layer,
+                                   downsample=None))
+
+        return nn.Sequential(*blocks)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        features = [x]
+        for i in range(self.num_stages):
+            x = self.stages[i](x)
+            features.append(x)
+        features = features[-self.num_heads:]
+
+        if self.use_fpn:
+            features = self.fpn(features)
+        if self.use_sam:
+            features = self.sam(features)
+        else:
+            final_size = features[-1].shape[2:]
+            for i in range(self.num_heads - 1):
+                features[i] = F.adaptive_avg_pool2d(features[i], final_size)
+
+        outputs = []
+        for head_idx in range(self.num_heads):
+            codes = self.heads[head_idx](features[head_idx])
+            start_latent_idx = sum(self.num_latents_per_head[:head_idx])
+            end_latent_idx = sum(self.num_latents_per_head[:head_idx + 1])
+            split_size = self.latent_dim[start_latent_idx:end_latent_idx]
+            outputs.extend(torch.split(codes, split_size, dim=1))
+        max_dim = max(self.latent_dim)
+        for i, dim in enumerate(self.latent_dim):
+            if dim < max_dim:
+                outputs[i] = F.pad(outputs[i], (0, max_dim - dim))
+            outputs[i] = outputs[i].unsqueeze(1)
+
+        return torch.cat(outputs, dim=1)
+
+
+class FPN(nn.Module):
+    """Implementation of Feature Pyramid Network (FPN).
+
+    The input of this module is a pyramid of features with reducing resolutions.
+    Then, this module fuses these multi-level features from `top_level` to
+    `bottom_level`. In particular, starting from the `top_level`, each feature
+    is convoluted, upsampled, and fused into its previous feature (which is also
+    convoluted).
+
+    Args:
+        pyramid_channels: A list of integers, each of which indicates the number
+            of channels of the feature from a particular level.
+        out_channels: Number of channels for each output.
+
+    Returns:
+        A list of feature maps, each of which has `out_channels` channels.
+    """
+
+    def __init__(self, pyramid_channels, out_channels):
+        super().__init__()
+        assert isinstance(pyramid_channels, (list, tuple))
+        self.num_levels = len(pyramid_channels)
+
+        self.lateral_layers = nn.ModuleList()
+        self.feature_layers = nn.ModuleList()
+        for i in range(self.num_levels):
+            in_channels = pyramid_channels[i]
+            self.lateral_layers.append(nn.Conv2d(in_channels=in_channels,
+                                                 out_channels=out_channels,
+                                                 kernel_size=3,
+                                                 padding=1,
+                                                 bias=True))
+            self.feature_layers.append(nn.Conv2d(in_channels=out_channels,
+                                                 out_channels=out_channels,
+                                                 kernel_size=3,
+                                                 padding=1,
+                                                 bias=True))
+
+    def forward(self, inputs):
+        if len(inputs) != self.num_levels:
+            raise ValueError('Number of inputs and `num_levels` mismatch!')
+
+        # Project all related features to `out_channels`.
+        laterals = []
+        for i in range(self.num_levels):
+            laterals.append(self.lateral_layers[i](inputs[i]))
+
+        # Fusion, starting from `top_level`.
+        for i in range(self.num_levels - 1, 0, -1):
+            scale_factor = laterals[i - 1].shape[2] // laterals[i].shape[2]
+            laterals[i - 1] = (laterals[i - 1] +
+                               F.interpolate(laterals[i],
+                                             mode='nearest',
+                                             scale_factor=scale_factor))
+
+        # Get outputs.
+        outputs = []
+        for i, lateral in enumerate(laterals):
+            outputs.append(self.feature_layers[i](lateral))
+
+        return outputs
+
+
+class SAM(nn.Module):
+    """Implementation of Spatial Alignment Module (SAM).
+
+    The input of this module is a pyramid of features with reducing resolutions.
+    Then this module downsamples all levels of feature to the minimum resolution
+    and fuses it with the smallest feature map.
+
+    Args:
+        pyramid_channels: A list of integers, each of which indicates the number
+            of channels of the feature from a particular level.
+        out_channels: Number of channels for each output.
+
+    Returns:
+        A list of feature maps, each of which has `out_channels` channels.
+    """
+
+    def __init__(self, pyramid_channels, out_channels):
+        super().__init__()
+        assert isinstance(pyramid_channels, (list, tuple))
+        self.num_levels = len(pyramid_channels)
+
+        self.fusion_layers = nn.ModuleList()
+        for i in range(self.num_levels):
+            in_channels = pyramid_channels[i]
+            self.fusion_layers.append(nn.Conv2d(in_channels=in_channels,
+                                                out_channels=out_channels,
+                                                kernel_size=3,
+                                                padding=1,
+                                                bias=True))
+
+    def forward(self, inputs):
+        if len(inputs) != self.num_levels:
+            raise ValueError('Number of inputs and `num_levels` mismatch!')
+
+        output_res = inputs[-1].shape[2:]
+        for i in range(self.num_levels - 1, -1, -1):
+            if i != self.num_levels - 1:
+                inputs[i] = F.adaptive_avg_pool2d(inputs[i], output_res)
+            inputs[i] = self.fusion_layers[i](inputs[i])
+            if i != self.num_levels - 1:
+                inputs[i] = inputs[i] + inputs[-1]
+
+        return inputs
+
+
+class CodeHead(nn.Module):
+    """Implementation of the task-head to produce inverted codes."""
+
+    def __init__(self, in_channels, out_channels, norm_layer):
+        super().__init__()
+        self.fc = nn.Linear(in_channels, out_channels, bias=True)
+        if norm_layer is None:
+            self.norm = nn.Identity()
+        else:
+            self.norm = norm_layer(out_channels)
+
+    def forward(self, x):
+        if x.ndim > 2:
+            x = x.flatten(start_dim=1)
+        latent = self.fc(x)
+        latent = latent.unsqueeze(2).unsqueeze(3)
+        latent = self.norm(latent)
+
+        return latent.flatten(start_dim=1)
+
+# pylint: enable=missing-function-docstring

models/inception_model.py

@@ -0,0 +1,562 @@
+# python3.7
+"""Contains the Inception V3 model, which is used for inference ONLY.
+
+This file is mostly borrowed from `torchvision/models/inception.py`.
+
+Inception model is widely used to compute FID or IS metric for evaluating
+generative models. However, the pre-trained models from torchvision is slightly
+different from the TensorFlow version
+
+http://download.tensorflow.org/models/image/imagenet/inception-2015-12-05.tgz
+
+which is used by the official FID implementation
+
+https://github.com/bioinf-jku/TTUR
+
+In particular:
+
+(1) The number of classes in TensorFlow model is 1008 instead of 1000.
+(2) The avg_pool() layers in TensorFlow model does not include the padded zero.
+(3) The last Inception E Block in TensorFlow model use max_pool() instead of
+    avg_pool().
+
+Hence, to align the evaluation results with those from TensorFlow
+implementation, we modified the inception model to support both versions. Please
+use `align_tf` argument to control the version.
+"""
+
+import warnings
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.distributed as dist
+
+from utils.misc import download_url
+
+__all__ = ['InceptionModel']
+
+# pylint: disable=line-too-long
+
+_MODEL_URL_SHA256 = {
+    # This model is provided by `torchvision`, which is ported from TensorFlow.
+    'torchvision_official': (
+        'https://download.pytorch.org/models/inception_v3_google-1a9a5a14.pth',
+        '1a9a5a14f40645a370184bd54f4e8e631351e71399112b43ad0294a79da290c8'  # hash sha256
+    ),
+
+    # This model is provided by https://github.com/mseitzer/pytorch-fid
+    'tf_inception_v3': (
+        'https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth',
+        '6726825d0af5f729cebd5821db510b11b1cfad8faad88a03f1befd49fb9129b2'  # hash sha256
+    )
+}
+
+
+class InceptionModel(object):
+    """Defines the Inception (V3) model.
+
+    This is a static class, which is used to avoid this model to be built
+    repeatedly. Consequently, this model is particularly used for inference,
+    like computing FID. If training is required, please use the model from
+    `torchvision.models` or implement by yourself.
+
+    NOTE: The pre-trained model assumes the inputs to be with `RGB` channel
+    order and pixel range [-1, 1], and will also resize the images to shape
+    [299, 299] automatically. If your input is normalized by subtracting
+    (0.485, 0.456, 0.406) and dividing (0.229, 0.224, 0.225), please use
+    `transform_input` in the `forward()` function to un-normalize it.
+    """
+    models = dict()
+
+    @staticmethod
+    def build_model(align_tf=True):
+        """Builds the model and load pre-trained weights.
+
+        If `align_tf` is set as True, the model will predict 1008 classes, and
+        the pre-trained weight from `https://github.com/mseitzer/pytorch-fid`
+        will be loaded. Otherwise, the model will predict 1000 classes, and will
+        load the model from `torchvision`.
+
+        The built model supports following arguments when forwarding:
+
+        - transform_input: Whether to transform the input back to pixel range
+            (-1, 1). Please disable this argument if your input is already with
+            pixel range (-1, 1). (default: False)
+        - output_logits: Whether to output the categorical logits instead of
+            features. (default: False)
+        - remove_logits_bias: Whether to remove the bias when computing the
+            logits. The official implementation removes the bias by default.
+            Please refer to
+            `https://github.com/openai/improved-gan/blob/master/inception_score/model.py`.
+            (default: False)
+        - output_predictions: Whether to output the final predictions, i.e.,
+            `softmax(logits)`. (default: False)
+        """
+        if align_tf:
+            num_classes = 1008
+            model_source = 'tf_inception_v3'
+        else:
+            num_classes = 1000
+            model_source = 'torchvision_official'
+
+        fingerprint = model_source
+
+        if fingerprint not in InceptionModel.models:
+            # Build model.
+            model = Inception3(num_classes=num_classes,
+                               aux_logits=False,
+                               init_weights=False,
+                               align_tf=align_tf)
+
+            # Download pre-trained weights.
+            if dist.is_initialized() and dist.get_rank() != 0:
+                dist.barrier()  # Download by chief.
+
+            url, sha256 = _MODEL_URL_SHA256[model_source]
+            filename = f'inception_model_{model_source}_{sha256}.pth'
+            model_path, hash_check = download_url(url,
+                                                  filename=filename,
+                                                  sha256=sha256)
+            state_dict = torch.load(model_path, map_location='cpu')
+            if hash_check is False:
+                warnings.warn(f'Hash check failed! The remote file from URL '
+                              f'`{url}` may be changed, or the downloading is '
+                              f'interrupted. The loaded inception model may '
+                              f'have unexpected behavior.')
+
+            if dist.is_initialized() and dist.get_rank() == 0:
+                dist.barrier()  # Wait for other replicas.
+
+            # Load weights.
+            model.load_state_dict(state_dict, strict=False)
+            del state_dict
+
+            # For inference only.
+            model.eval().requires_grad_(False).cuda()
+            InceptionModel.models[fingerprint] = model
+
+        return InceptionModel.models[fingerprint]
+
+# pylint: disable=missing-function-docstring
+# pylint: disable=missing-class-docstring
+# pylint: disable=super-with-arguments
+# pylint: disable=consider-merging-isinstance
+# pylint: disable=import-outside-toplevel
+# pylint: disable=no-else-return
+
+class Inception3(nn.Module):
+
+    def __init__(self, num_classes=1000, aux_logits=True, inception_blocks=None,
+                 init_weights=True, align_tf=True):
+        super(Inception3, self).__init__()
+        if inception_blocks is None:
+            inception_blocks = [
+                BasicConv2d, InceptionA, InceptionB, InceptionC,
+                InceptionD, InceptionE, InceptionAux
+            ]
+        assert len(inception_blocks) == 7
+        conv_block = inception_blocks[0]
+        inception_a = inception_blocks[1]
+        inception_b = inception_blocks[2]
+        inception_c = inception_blocks[3]
+        inception_d = inception_blocks[4]
+        inception_e = inception_blocks[5]
+        inception_aux = inception_blocks[6]
+
+        self.aux_logits = aux_logits
+        self.align_tf = align_tf
+        self.Conv2d_1a_3x3 = conv_block(3, 32, kernel_size=3, stride=2)
+        self.Conv2d_2a_3x3 = conv_block(32, 32, kernel_size=3)
+        self.Conv2d_2b_3x3 = conv_block(32, 64, kernel_size=3, padding=1)
+        self.Conv2d_3b_1x1 = conv_block(64, 80, kernel_size=1)
+        self.Conv2d_4a_3x3 = conv_block(80, 192, kernel_size=3)
+        self.Mixed_5b = inception_a(192, pool_features=32, align_tf=self.align_tf)
+        self.Mixed_5c = inception_a(256, pool_features=64, align_tf=self.align_tf)
+        self.Mixed_5d = inception_a(288, pool_features=64, align_tf=self.align_tf)
+        self.Mixed_6a = inception_b(288)
+        self.Mixed_6b = inception_c(768, channels_7x7=128, align_tf=self.align_tf)
+        self.Mixed_6c = inception_c(768, channels_7x7=160, align_tf=self.align_tf)
+        self.Mixed_6d = inception_c(768, channels_7x7=160, align_tf=self.align_tf)
+        self.Mixed_6e = inception_c(768, channels_7x7=192, align_tf=self.align_tf)
+        if aux_logits:
+            self.AuxLogits = inception_aux(768, num_classes)
+        self.Mixed_7a = inception_d(768)
+        self.Mixed_7b = inception_e(1280, align_tf=self.align_tf)
+        self.Mixed_7c = inception_e(2048, use_max_pool=self.align_tf)
+        self.fc = nn.Linear(2048, num_classes)
+        if init_weights:
+            for m in self.modules():
+                if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
+                    import scipy.stats as stats
+                    stddev = m.stddev if hasattr(m, 'stddev') else 0.1
+                    X = stats.truncnorm(-2, 2, scale=stddev)
+                    values = torch.as_tensor(X.rvs(m.weight.numel()), dtype=m.weight.dtype)
+                    values = values.view(m.weight.size())
+                    with torch.no_grad():
+                        m.weight.copy_(values)
+                elif isinstance(m, nn.BatchNorm2d):
+                    nn.init.constant_(m.weight, 1)
+                    nn.init.constant_(m.bias, 0)
+
+    @staticmethod
+    def _transform_input(x, transform_input=False):
+        if transform_input:
+            x_ch0 = torch.unsqueeze(x[:, 0], 1) * (0.229 / 0.5) + (0.485 - 0.5) / 0.5
+            x_ch1 = torch.unsqueeze(x[:, 1], 1) * (0.224 / 0.5) + (0.456 - 0.5) / 0.5
+            x_ch2 = torch.unsqueeze(x[:, 2], 1) * (0.225 / 0.5) + (0.406 - 0.5) / 0.5
+            x = torch.cat((x_ch0, x_ch1, x_ch2), 1)
+        return x
+
+    def _forward(self,
+                 x,
+                 output_logits=False,
+                 remove_logits_bias=False,
+                 output_predictions=False):
+        # Upsample if necessary.
+        if x.shape[2] != 299 or x.shape[3] != 299:
+            if self.align_tf:
+                theta = torch.eye(2, 3).to(x)
+                theta[0, 2] += theta[0, 0] / x.shape[3] - theta[0, 0] / 299
+                theta[1, 2] += theta[1, 1] / x.shape[2] - theta[1, 1] / 299
+                theta = theta.unsqueeze(0).repeat(x.shape[0], 1, 1)
+                grid = F.affine_grid(theta,
+                                     size=(x.shape[0], x.shape[1], 299, 299),
+                                     align_corners=False)
+                x = F.grid_sample(x, grid,
+                                  mode='bilinear',
+                                  padding_mode='border',
+                                  align_corners=False)
+            else:
+                x = F.interpolate(
+                    x, size=(299, 299), mode='bilinear', align_corners=False)
+        if x.shape[1] == 1:
+            x = x.repeat((1, 3, 1, 1))
+
+        if self.align_tf:
+            x = (x * 127.5 + 127.5 - 128) / 128
+
+        # N x 3 x 299 x 299
+        x = self.Conv2d_1a_3x3(x)
+        # N x 32 x 149 x 149
+        x = self.Conv2d_2a_3x3(x)
+        # N x 32 x 147 x 147
+        x = self.Conv2d_2b_3x3(x)
+        # N x 64 x 147 x 147
+        x = F.max_pool2d(x, kernel_size=3, stride=2)
+        # N x 64 x 73 x 73
+        x = self.Conv2d_3b_1x1(x)
+        # N x 80 x 73 x 73
+        x = self.Conv2d_4a_3x3(x)
+        # N x 192 x 71 x 71
+        x = F.max_pool2d(x, kernel_size=3, stride=2)
+        # N x 192 x 35 x 35
+        x = self.Mixed_5b(x)
+        # N x 256 x 35 x 35
+        x = self.Mixed_5c(x)
+        # N x 288 x 35 x 35
+        x = self.Mixed_5d(x)
+        # N x 288 x 35 x 35
+        x = self.Mixed_6a(x)
+        # N x 768 x 17 x 17
+        x = self.Mixed_6b(x)
+        # N x 768 x 17 x 17
+        x = self.Mixed_6c(x)
+        # N x 768 x 17 x 17
+        x = self.Mixed_6d(x)
+        # N x 768 x 17 x 17
+        x = self.Mixed_6e(x)
+        # N x 768 x 17 x 17
+        if self.training and self.aux_logits:
+            aux = self.AuxLogits(x)
+        else:
+            aux = None
+        # N x 768 x 17 x 17
+        x = self.Mixed_7a(x)
+        # N x 1280 x 8 x 8
+        x = self.Mixed_7b(x)
+        # N x 2048 x 8 x 8
+        x = self.Mixed_7c(x)
+        # N x 2048 x 8 x 8
+        # Adaptive average pooling
+        x = F.adaptive_avg_pool2d(x, (1, 1))
+        # N x 2048 x 1 x 1
+        x = F.dropout(x, training=self.training)
+        # N x 2048 x 1 x 1
+        x = torch.flatten(x, 1)
+        # N x 2048
+        if output_logits or output_predictions:
+            x = self.fc(x)
+            # N x 1000 (num_classes)
+            if remove_logits_bias:
+                x = x - self.fc.bias.view(1, -1)
+            if output_predictions:
+                x = F.softmax(x, dim=1)
+        return x, aux
+
+    def forward(self,
+                x,
+                transform_input=False,
+                output_logits=False,
+                remove_logits_bias=False,
+                output_predictions=False):
+        x = self._transform_input(x, transform_input)
+        x, aux = self._forward(
+            x, output_logits, remove_logits_bias, output_predictions)
+        if self.training and self.aux_logits:
+            return x, aux
+        else:
+            return x
+
+
+class InceptionA(nn.Module):
+
+    def __init__(self, in_channels, pool_features, conv_block=None, align_tf=False):
+        super(InceptionA, self).__init__()
+        if conv_block is None:
+            conv_block = BasicConv2d
+        self.branch1x1 = conv_block(in_channels, 64, kernel_size=1)
+
+        self.branch5x5_1 = conv_block(in_channels, 48, kernel_size=1)
+        self.branch5x5_2 = conv_block(48, 64, kernel_size=5, padding=2)
+
+        self.branch3x3dbl_1 = conv_block(in_channels, 64, kernel_size=1)
+        self.branch3x3dbl_2 = conv_block(64, 96, kernel_size=3, padding=1)
+        self.branch3x3dbl_3 = conv_block(96, 96, kernel_size=3, padding=1)
+
+        self.branch_pool = conv_block(in_channels, pool_features, kernel_size=1)
+        self.pool_include_padding = not align_tf
+
+    def _forward(self, x):
+        branch1x1 = self.branch1x1(x)
+
+        branch5x5 = self.branch5x5_1(x)
+        branch5x5 = self.branch5x5_2(branch5x5)
+
+        branch3x3dbl = self.branch3x3dbl_1(x)
+        branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
+        branch3x3dbl = self.branch3x3dbl_3(branch3x3dbl)
+
+        branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1,
+                                   count_include_pad=self.pool_include_padding)
+        branch_pool = self.branch_pool(branch_pool)
+
+        outputs = [branch1x1, branch5x5, branch3x3dbl, branch_pool]
+        return outputs
+
+    def forward(self, x):
+        outputs = self._forward(x)
+        return torch.cat(outputs, 1)
+
+
+class InceptionB(nn.Module):
+
+    def __init__(self, in_channels, conv_block=None):
+        super(InceptionB, self).__init__()
+        if conv_block is None:
+            conv_block = BasicConv2d
+        self.branch3x3 = conv_block(in_channels, 384, kernel_size=3, stride=2)
+
+        self.branch3x3dbl_1 = conv_block(in_channels, 64, kernel_size=1)
+        self.branch3x3dbl_2 = conv_block(64, 96, kernel_size=3, padding=1)
+        self.branch3x3dbl_3 = conv_block(96, 96, kernel_size=3, stride=2)
+
+    def _forward(self, x):
+        branch3x3 = self.branch3x3(x)
+
+        branch3x3dbl = self.branch3x3dbl_1(x)
+        branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
+        branch3x3dbl = self.branch3x3dbl_3(branch3x3dbl)
+
+        branch_pool = F.max_pool2d(x, kernel_size=3, stride=2)
+
+        outputs = [branch3x3, branch3x3dbl, branch_pool]
+        return outputs
+
+    def forward(self, x):
+        outputs = self._forward(x)
+        return torch.cat(outputs, 1)
+
+
+class InceptionC(nn.Module):
+
+    def __init__(self, in_channels, channels_7x7, conv_block=None, align_tf=False):
+        super(InceptionC, self).__init__()
+        if conv_block is None:
+            conv_block = BasicConv2d
+        self.branch1x1 = conv_block(in_channels, 192, kernel_size=1)
+
+        c7 = channels_7x7
+        self.branch7x7_1 = conv_block(in_channels, c7, kernel_size=1)
+        self.branch7x7_2 = conv_block(c7, c7, kernel_size=(1, 7), padding=(0, 3))
+        self.branch7x7_3 = conv_block(c7, 192, kernel_size=(7, 1), padding=(3, 0))
+
+        self.branch7x7dbl_1 = conv_block(in_channels, c7, kernel_size=1)
+        self.branch7x7dbl_2 = conv_block(c7, c7, kernel_size=(7, 1), padding=(3, 0))
+        self.branch7x7dbl_3 = conv_block(c7, c7, kernel_size=(1, 7), padding=(0, 3))
+        self.branch7x7dbl_4 = conv_block(c7, c7, kernel_size=(7, 1), padding=(3, 0))
+        self.branch7x7dbl_5 = conv_block(c7, 192, kernel_size=(1, 7), padding=(0, 3))
+
+        self.branch_pool = conv_block(in_channels, 192, kernel_size=1)
+        self.pool_include_padding = not align_tf
+
+    def _forward(self, x):
+        branch1x1 = self.branch1x1(x)
+
+        branch7x7 = self.branch7x7_1(x)
+        branch7x7 = self.branch7x7_2(branch7x7)
+        branch7x7 = self.branch7x7_3(branch7x7)
+
+        branch7x7dbl = self.branch7x7dbl_1(x)
+        branch7x7dbl = self.branch7x7dbl_2(branch7x7dbl)
+        branch7x7dbl = self.branch7x7dbl_3(branch7x7dbl)
+        branch7x7dbl = self.branch7x7dbl_4(branch7x7dbl)
+        branch7x7dbl = self.branch7x7dbl_5(branch7x7dbl)
+
+        branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1,
+                                   count_include_pad=self.pool_include_padding)
+        branch_pool = self.branch_pool(branch_pool)
+
+        outputs = [branch1x1, branch7x7, branch7x7dbl, branch_pool]
+        return outputs
+
+    def forward(self, x):
+        outputs = self._forward(x)
+        return torch.cat(outputs, 1)
+
+
+class InceptionD(nn.Module):
+
+    def __init__(self, in_channels, conv_block=None):
+        super(InceptionD, self).__init__()
+        if conv_block is None:
+            conv_block = BasicConv2d
+        self.branch3x3_1 = conv_block(in_channels, 192, kernel_size=1)
+        self.branch3x3_2 = conv_block(192, 320, kernel_size=3, stride=2)
+
+        self.branch7x7x3_1 = conv_block(in_channels, 192, kernel_size=1)
+        self.branch7x7x3_2 = conv_block(192, 192, kernel_size=(1, 7), padding=(0, 3))
+        self.branch7x7x3_3 = conv_block(192, 192, kernel_size=(7, 1), padding=(3, 0))
+        self.branch7x7x3_4 = conv_block(192, 192, kernel_size=3, stride=2)
+
+    def _forward(self, x):
+        branch3x3 = self.branch3x3_1(x)
+        branch3x3 = self.branch3x3_2(branch3x3)
+
+        branch7x7x3 = self.branch7x7x3_1(x)
+        branch7x7x3 = self.branch7x7x3_2(branch7x7x3)
+        branch7x7x3 = self.branch7x7x3_3(branch7x7x3)
+        branch7x7x3 = self.branch7x7x3_4(branch7x7x3)
+
+        branch_pool = F.max_pool2d(x, kernel_size=3, stride=2)
+        outputs = [branch3x3, branch7x7x3, branch_pool]
+        return outputs
+
+    def forward(self, x):
+        outputs = self._forward(x)
+        return torch.cat(outputs, 1)
+
+
+class InceptionE(nn.Module):
+
+    def __init__(self, in_channels, conv_block=None, align_tf=False, use_max_pool=False):
+        super(InceptionE, self).__init__()
+        if conv_block is None:
+            conv_block = BasicConv2d
+        self.branch1x1 = conv_block(in_channels, 320, kernel_size=1)
+
+        self.branch3x3_1 = conv_block(in_channels, 384, kernel_size=1)
+        self.branch3x3_2a = conv_block(384, 384, kernel_size=(1, 3), padding=(0, 1))
+        self.branch3x3_2b = conv_block(384, 384, kernel_size=(3, 1), padding=(1, 0))
+
+        self.branch3x3dbl_1 = conv_block(in_channels, 448, kernel_size=1)
+        self.branch3x3dbl_2 = conv_block(448, 384, kernel_size=3, padding=1)
+        self.branch3x3dbl_3a = conv_block(384, 384, kernel_size=(1, 3), padding=(0, 1))
+        self.branch3x3dbl_3b = conv_block(384, 384, kernel_size=(3, 1), padding=(1, 0))
+
+        self.branch_pool = conv_block(in_channels, 192, kernel_size=1)
+        self.pool_include_padding = not align_tf
+        self.use_max_pool = use_max_pool
+
+    def _forward(self, x):
+        branch1x1 = self.branch1x1(x)
+
+        branch3x3 = self.branch3x3_1(x)
+        branch3x3 = [
+            self.branch3x3_2a(branch3x3),
+            self.branch3x3_2b(branch3x3),
+        ]
+        branch3x3 = torch.cat(branch3x3, 1)
+
+        branch3x3dbl = self.branch3x3dbl_1(x)
+        branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
+        branch3x3dbl = [
+            self.branch3x3dbl_3a(branch3x3dbl),
+            self.branch3x3dbl_3b(branch3x3dbl),
+        ]
+        branch3x3dbl = torch.cat(branch3x3dbl, 1)
+
+        if self.use_max_pool:
+            branch_pool = F.max_pool2d(x, kernel_size=3, stride=1, padding=1)
+        else:
+            branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1,
+                                       count_include_pad=self.pool_include_padding)
+        branch_pool = self.branch_pool(branch_pool)
+
+        outputs = [branch1x1, branch3x3, branch3x3dbl, branch_pool]
+        return outputs
+
+    def forward(self, x):
+        outputs = self._forward(x)
+        return torch.cat(outputs, 1)
+
+
+class InceptionAux(nn.Module):
+
+    def __init__(self, in_channels, num_classes, conv_block=None):
+        super(InceptionAux, self).__init__()
+        if conv_block is None:
+            conv_block = BasicConv2d
+        self.conv0 = conv_block(in_channels, 128, kernel_size=1)
+        self.conv1 = conv_block(128, 768, kernel_size=5)
+        self.conv1.stddev = 0.01
+        self.fc = nn.Linear(768, num_classes)
+        self.fc.stddev = 0.001
+
+    def forward(self, x):
+        # N x 768 x 17 x 17
+        x = F.avg_pool2d(x, kernel_size=5, stride=3)
+        # N x 768 x 5 x 5
+        x = self.conv0(x)
+        # N x 128 x 5 x 5
+        x = self.conv1(x)
+        # N x 768 x 1 x 1
+        # Adaptive average pooling
+        x = F.adaptive_avg_pool2d(x, (1, 1))
+        # N x 768 x 1 x 1
+        x = torch.flatten(x, 1)
+        # N x 768
+        x = self.fc(x)
+        # N x 1000
+        return x
+
+
+class BasicConv2d(nn.Module):
+
+    def __init__(self, in_channels, out_channels, **kwargs):
+        super(BasicConv2d, self).__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, bias=False, **kwargs)
+        self.bn = nn.BatchNorm2d(out_channels, eps=0.001)
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        return F.relu(x, inplace=True)
+
+# pylint: enable=line-too-long
+# pylint: enable=missing-function-docstring
+# pylint: enable=missing-class-docstring
+# pylint: enable=super-with-arguments
+# pylint: enable=consider-merging-isinstance
+# pylint: enable=import-outside-toplevel
+# pylint: enable=no-else-return

models/perceptual_model.py

@@ -0,0 +1,519 @@
+# python3.7
+"""Contains the VGG16 model, which is used for inference ONLY.
+
+VGG16 is commonly used for perceptual feature extraction. The model implemented
+in this file can be used for evaluation (like computing LPIPS, perceptual path
+length, etc.), OR be used in training for loss computation (like perceptual
+loss, etc.).
+
+The pre-trained model is officially shared by
+
+https://www.robots.ox.ac.uk/~vgg/research/very_deep/
+
+and ported by
+
+https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metrics/vgg16.pt
+
+Compared to the official VGG16 model, this ported model also support evaluating
+LPIPS, which is introduced in
+
+https://github.com/richzhang/PerceptualSimilarity
+"""
+
+import warnings
+import numpy as np
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.distributed as dist
+
+from utils.misc import download_url
+
+__all__ = ['PerceptualModel']
+
+# pylint: disable=line-too-long
+_MODEL_URL_SHA256 = {
+    # This model is provided by `torchvision`, which is ported from TensorFlow.
+    'torchvision_official': (
+        'https://download.pytorch.org/models/vgg16-397923af.pth',
+        '397923af8e79cdbb6a7127f12361acd7a2f83e06b05044ddf496e83de57a5bf0'  # hash sha256
+    ),
+
+    # This model is provided by https://github.com/NVlabs/stylegan2-ada-pytorch
+    'vgg_perceptual_lpips': (
+        'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metrics/vgg16.pt',
+        'b437eb095feaeb0b83eb3fa11200ebca4548ee39a07fb944a417ddc516cc07c3'  # hash sha256
+    )
+}
+# pylint: enable=line-too-long
+
+
+class PerceptualModel(object):
+    """Defines the perceptual model, which is based on VGG16 structure.
+
+    This is a static class, which is used to avoid this model to be built
+    repeatedly. Consequently, this model is particularly used for inference,
+    like computing LPIPS, or for loss computation, like perceptual loss. If
+    training is required, please use the model from `torchvision.models` or
+    implement by yourself.
+
+    NOTE: The pre-trained model assumes the inputs to be with `RGB` channel
+    order and pixel range [-1, 1], and will NOT resize the input automatically
+    if only perceptual feature is needed.
+    """
+    models = dict()
+
+    @staticmethod
+    def build_model(use_torchvision=False, no_top=True, enable_lpips=True):
+        """Builds the model and load pre-trained weights.
+
+        1. If `use_torchvision` is set as True, the model released by
+           `torchvision` will be loaded, otherwise, the model released by
+           https://www.robots.ox.ac.uk/~vgg/research/very_deep/ will be used.
+           (default: False)
+
+        2. To save computing resources, these is an option to only load the
+           backbone (i.e., without the last three fully-connected layers). This
+           is commonly used for perceptual loss or LPIPS loss computation.
+           Please use argument `no_top` to control this. (default: True)
+
+        3. For LPIPS loss computation, some additional weights (which is used
+           for balancing the features from different resolutions) are employed
+           on top of the original VGG16 backbone. Details can be found at
+           https://github.com/richzhang/PerceptualSimilarity. Please use
+           `enable_lpips` to enable this feature. (default: True)
+
+        The built model supports following arguments when forwarding:
+
+        - resize_input: Whether to resize the input image to size [224, 224]
+            before forwarding. For feature-based computation (i.e., only
+            convolutional layers are used), image resizing is not essential.
+            (default: False)
+        - return_tensor: This field resolves the model behavior. Following
+            options are supported:
+                `feature1`: Before the first max pooling layer.
+                `pool1`: After the first max pooling layer.
+                `feature2`: Before the second max pooling layer.
+                `pool2`: After the second max pooling layer.
+                `feature3`: Before the third max pooling layer.
+                `pool3`: After the third max pooling layer.
+                `feature4`: Before the fourth max pooling layer.
+                `pool4`: After the fourth max pooling layer.
+                `feature5`: Before the fifth max pooling layer.
+                `pool5`: After the fifth max pooling layer.
+                `flatten`: The flattened feature, after `adaptive_avgpool`.
+                `feature`: The 4096d feature for logits computation. (default)
+                `logits`: The 1000d categorical logits.
+                `prediction`: The 1000d predicted probability.
+                `lpips`: The LPIPS score between two input images.
+        """
+        if use_torchvision:
+            model_source = 'torchvision_official'
+            align_tf_resize = False
+            is_torch_script = False
+        else:
+            model_source = 'vgg_perceptual_lpips'
+            align_tf_resize = True
+            is_torch_script = True
+
+        if enable_lpips and model_source != 'vgg_perceptual_lpips':
+            warnings.warn('The pre-trained model officially released by '
+                          '`torchvision` does not support LPIPS computation! '
+                          'Equal weights will be used for each resolution.')
+
+        fingerprint = (model_source, no_top, enable_lpips)
+
+        if fingerprint not in PerceptualModel.models:
+            # Build model.
+            model = VGG16(align_tf_resize=align_tf_resize,
+                          no_top=no_top,
+                          enable_lpips=enable_lpips)
+
+            # Download pre-trained weights.
+            if dist.is_initialized() and dist.get_rank() != 0:
+                dist.barrier()  # Download by chief.
+
+            url, sha256 = _MODEL_URL_SHA256[model_source]
+            filename = f'perceptual_model_{model_source}_{sha256}.pth'
+            model_path, hash_check = download_url(url,
+                                                  filename=filename,
+                                                  sha256=sha256)
+            if is_torch_script:
+                src_state_dict = torch.jit.load(model_path, map_location='cpu')
+            else:
+                src_state_dict = torch.load(model_path, map_location='cpu')
+            if hash_check is False:
+                warnings.warn(f'Hash check failed! The remote file from URL '
+                              f'`{url}` may be changed, or the downloading is '
+                              f'interrupted. The loaded perceptual model may '
+                              f'have unexpected behavior.')
+
+            if dist.is_initialized() and dist.get_rank() == 0:
+                dist.barrier()  # Wait for other replicas.
+
+            # Load weights.
+            dst_state_dict = _convert_weights(src_state_dict, model_source)
+            model.load_state_dict(dst_state_dict, strict=False)
+            del src_state_dict, dst_state_dict
+
+            # For inference only.
+            model.eval().requires_grad_(False).cuda()
+            PerceptualModel.models[fingerprint] = model
+
+        return PerceptualModel.models[fingerprint]
+
+
+def _convert_weights(src_state_dict, model_source):
+    if model_source not in _MODEL_URL_SHA256:
+        raise ValueError(f'Invalid model source `{model_source}`!\n'
+                         f'Sources allowed: {list(_MODEL_URL_SHA256.keys())}.')
+    if model_source == 'torchvision_official':
+        dst_to_src_var_mapping = {
+            'conv11.weight': 'features.0.weight',
+            'conv11.bias': 'features.0.bias',
+            'conv12.weight': 'features.2.weight',
+            'conv12.bias': 'features.2.bias',
+            'conv21.weight': 'features.5.weight',
+            'conv21.bias': 'features.5.bias',
+            'conv22.weight': 'features.7.weight',
+            'conv22.bias': 'features.7.bias',
+            'conv31.weight': 'features.10.weight',
+            'conv31.bias': 'features.10.bias',
+            'conv32.weight': 'features.12.weight',
+            'conv32.bias': 'features.12.bias',
+            'conv33.weight': 'features.14.weight',
+            'conv33.bias': 'features.14.bias',
+            'conv41.weight': 'features.17.weight',
+            'conv41.bias': 'features.17.bias',
+            'conv42.weight': 'features.19.weight',
+            'conv42.bias': 'features.19.bias',
+            'conv43.weight': 'features.21.weight',
+            'conv43.bias': 'features.21.bias',
+            'conv51.weight': 'features.24.weight',
+            'conv51.bias': 'features.24.bias',
+            'conv52.weight': 'features.26.weight',
+            'conv52.bias': 'features.26.bias',
+            'conv53.weight': 'features.28.weight',
+            'conv53.bias': 'features.28.bias',
+            'fc1.weight': 'classifier.0.weight',
+            'fc1.bias': 'classifier.0.bias',
+            'fc2.weight': 'classifier.3.weight',
+            'fc2.bias': 'classifier.3.bias',
+            'fc3.weight': 'classifier.6.weight',
+            'fc3.bias': 'classifier.6.bias',
+        }
+    elif model_source == 'vgg_perceptual_lpips':
+        src_state_dict = src_state_dict.state_dict()
+        dst_to_src_var_mapping = {
+            'conv11.weight': 'layers.conv1.weight',
+            'conv11.bias': 'layers.conv1.bias',
+            'conv12.weight': 'layers.conv2.weight',
+            'conv12.bias': 'layers.conv2.bias',
+            'conv21.weight': 'layers.conv3.weight',
+            'conv21.bias': 'layers.conv3.bias',
+            'conv22.weight': 'layers.conv4.weight',
+            'conv22.bias': 'layers.conv4.bias',
+            'conv31.weight': 'layers.conv5.weight',
+            'conv31.bias': 'layers.conv5.bias',
+            'conv32.weight': 'layers.conv6.weight',
+            'conv32.bias': 'layers.conv6.bias',
+            'conv33.weight': 'layers.conv7.weight',
+            'conv33.bias': 'layers.conv7.bias',
+            'conv41.weight': 'layers.conv8.weight',
+            'conv41.bias': 'layers.conv8.bias',
+            'conv42.weight': 'layers.conv9.weight',
+            'conv42.bias': 'layers.conv9.bias',
+            'conv43.weight': 'layers.conv10.weight',
+            'conv43.bias': 'layers.conv10.bias',
+            'conv51.weight': 'layers.conv11.weight',
+            'conv51.bias': 'layers.conv11.bias',
+            'conv52.weight': 'layers.conv12.weight',
+            'conv52.bias': 'layers.conv12.bias',
+            'conv53.weight': 'layers.conv13.weight',
+            'conv53.bias': 'layers.conv13.bias',
+            'fc1.weight': 'layers.fc1.weight',
+            'fc1.bias': 'layers.fc1.bias',
+            'fc2.weight': 'layers.fc2.weight',
+            'fc2.bias': 'layers.fc2.bias',
+            'fc3.weight': 'layers.fc3.weight',
+            'fc3.bias': 'layers.fc3.bias',
+            'lpips.0.weight': 'lpips0',
+            'lpips.1.weight': 'lpips1',
+            'lpips.2.weight': 'lpips2',
+            'lpips.3.weight': 'lpips3',
+            'lpips.4.weight': 'lpips4',
+        }
+    else:
+        raise NotImplementedError(f'Not implemented model source '
+                                  f'`{model_source}`!')
+
+    dst_state_dict = {}
+    for dst_name, src_name in dst_to_src_var_mapping.items():
+        if dst_name.startswith('lpips'):
+            dst_state_dict[dst_name] = src_state_dict[src_name].unsqueeze(0)
+        else:
+            dst_state_dict[dst_name] = src_state_dict[src_name].clone()
+    return dst_state_dict
+
+
+_IMG_MEAN = (0.485, 0.456, 0.406)
+_IMG_STD  = (0.229, 0.224, 0.225)
+_ALLOWED_RETURN = [
+    'feature1', 'pool1', 'feature2', 'pool2', 'feature3', 'pool3', 'feature4',
+    'pool4', 'feature5', 'pool5', 'flatten', 'feature', 'logits', 'prediction',
+    'lpips'
+]
+
+# pylint: disable=missing-function-docstring
+
+class VGG16(nn.Module):
+    """Defines the VGG16 structure.
+
+    This model takes `RGB` images with data format `NCHW` as the raw inputs. The
+    pixel range are assumed to be [-1, 1].
+    """
+
+    def __init__(self, align_tf_resize=False, no_top=True, enable_lpips=True):
+        """Defines the network structure."""
+        super().__init__()
+
+        self.align_tf_resize = align_tf_resize
+        self.no_top = no_top
+        self.enable_lpips = enable_lpips
+
+        self.conv11 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1)
+        self.relu11 = nn.ReLU(inplace=True)
+        self.conv12 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)
+        self.relu12 = nn.ReLU(inplace=True)
+        # output `feature1`, with shape [N, 64, 224, 224]
+
+        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
+        # output `pool1`, with shape [N, 64, 112, 112]
+
+        self.conv21 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)
+        self.relu21 = nn.ReLU(inplace=True)
+        self.conv22 = nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1)
+        self.relu22 = nn.ReLU(inplace=True)
+        # output `feature2`, with shape [N, 128, 112, 112]
+
+        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
+        # output `pool2`, with shape [N, 128, 56, 56]
+
+        self.conv31 = nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1)
+        self.relu31 = nn.ReLU(inplace=True)
+        self.conv32 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
+        self.relu32 = nn.ReLU(inplace=True)
+        self.conv33 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
+        self.relu33 = nn.ReLU(inplace=True)
+        # output `feature3`, with shape [N, 256, 56, 56]
+
+        self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)
+        # output `pool3`, with shape [N,256, 28, 28]
+
+        self.conv41 = nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1)
+        self.relu41 = nn.ReLU(inplace=True)
+        self.conv42 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
+        self.relu42 = nn.ReLU(inplace=True)
+        self.conv43 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
+        self.relu43 = nn.ReLU(inplace=True)
+        # output `feature4`, with shape [N, 512, 28, 28]
+
+        self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2)
+        # output `pool4`, with shape [N, 512, 14, 14]
+
+        self.conv51 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
+        self.relu51 = nn.ReLU(inplace=True)
+        self.conv52 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
+        self.relu52 = nn.ReLU(inplace=True)
+        self.conv53 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
+        self.relu53 = nn.ReLU(inplace=True)
+        # output `feature5`, with shape [N, 512, 14, 14]
+
+        self.pool5 = nn.MaxPool2d(kernel_size=2, stride=2)
+        # output `pool5`, with shape [N, 512, 7, 7]
+
+        if self.enable_lpips:
+            self.lpips = nn.ModuleList()
+            for idx, ch in enumerate([64, 128, 256, 512, 512]):
+                self.lpips.append(nn.Conv2d(ch, 1, kernel_size=1, bias=False))
+                self.lpips[idx].weight.data.copy_(torch.ones(1, ch, 1, 1))
+
+        if not self.no_top:
+            self.avgpool = nn.AdaptiveAvgPool2d((7, 7))
+            self.flatten = nn.Flatten(start_dim=1, end_dim=-1)
+            # output `flatten`, with shape [N, 25088]
+
+            self.fc1 = nn.Linear(512 * 7 * 7, 4096)
+            self.fc1_relu = nn.ReLU(inplace=True)
+            self.fc1_dropout = nn.Dropout(0.5, inplace=False)
+            self.fc2 = nn.Linear(4096, 4096)
+            self.fc2_relu = nn.ReLU(inplace=True)
+            self.fc2_dropout = nn.Dropout(0.5, inplace=False)
+            # output `feature`, with shape [N, 4096]
+
+            self.fc3 = nn.Linear(4096, 1000)
+            # output `logits`, with shape [N, 1000]
+
+            self.out = nn.Softmax(dim=1)
+            # output `softmax`, with shape [N, 1000]
+
+        img_mean = np.array(_IMG_MEAN).reshape((1, 3, 1, 1)).astype(np.float32)
+        img_std = np.array(_IMG_STD).reshape((1, 3, 1, 1)).astype(np.float32)
+        self.register_buffer('img_mean', torch.from_numpy(img_mean))
+        self.register_buffer('img_std', torch.from_numpy(img_std))
+
+    def forward(self,
+                x,
+                y=None,
+                *,
+                resize_input=False,
+                return_tensor='feature'):
+        return_tensor = return_tensor.lower()
+        if return_tensor not in _ALLOWED_RETURN:
+            raise ValueError(f'Invalid output tensor name `{return_tensor}` '
+                             f'for perceptual model (VGG16)!\n'
+                             f'Names allowed: {_ALLOWED_RETURN}.')
+
+        if return_tensor == 'lpips' and y is None:
+            raise ValueError('Two images are required for LPIPS computation, '
+                             'but only one is received!')
+
+        if return_tensor == 'lpips':
+            assert x.shape == y.shape
+            x = torch.cat([x, y], dim=0)
+            features = []
+
+        if resize_input:
+            if self.align_tf_resize:
+                theta = torch.eye(2, 3).to(x)
+                theta[0, 2] += theta[0, 0] / x.shape[3] - theta[0, 0] / 224
+                theta[1, 2] += theta[1, 1] / x.shape[2] - theta[1, 1] / 224
+                theta = theta.unsqueeze(0).repeat(x.shape[0], 1, 1)
+                grid = F.affine_grid(theta,
+                                     size=(x.shape[0], x.shape[1], 224, 224),
+                                     align_corners=False)
+                x = F.grid_sample(x, grid,
+                                  mode='bilinear',
+                                  padding_mode='border',
+                                  align_corners=False)
+            else:
+                x = F.interpolate(x,
+                                  size=(224, 224),
+                                  mode='bilinear',
+                                  align_corners=False)
+        if x.shape[1] == 1:
+            x = x.repeat((1, 3, 1, 1))
+
+        x = (x + 1) / 2
+        x = (x - self.img_mean) / self.img_std
+
+        x = self.conv11(x)
+        x = self.relu11(x)
+        x = self.conv12(x)
+        x = self.relu12(x)
+        if return_tensor == 'feature1':
+            return x
+        if return_tensor == 'lpips':
+            features.append(x)
+
+        x = self.pool1(x)
+        if return_tensor == 'pool1':
+            return x
+
+        x = self.conv21(x)
+        x = self.relu21(x)
+        x = self.conv22(x)
+        x = self.relu22(x)
+        if return_tensor == 'feature2':
+            return x
+        if return_tensor == 'lpips':
+            features.append(x)
+
+        x = self.pool2(x)
+        if return_tensor == 'pool2':
+            return x
+
+        x = self.conv31(x)
+        x = self.relu31(x)
+        x = self.conv32(x)
+        x = self.relu32(x)
+        x = self.conv33(x)
+        x = self.relu33(x)
+        if return_tensor == 'feature3':
+            return x
+        if return_tensor == 'lpips':
+            features.append(x)
+
+        x = self.pool3(x)
+        if return_tensor == 'pool3':
+            return x
+
+        x = self.conv41(x)
+        x = self.relu41(x)
+        x = self.conv42(x)
+        x = self.relu42(x)
+        x = self.conv43(x)
+        x = self.relu43(x)
+        if return_tensor == 'feature4':
+            return x
+        if return_tensor == 'lpips':
+            features.append(x)
+
+        x = self.pool4(x)
+        if return_tensor == 'pool4':
+            return x
+
+        x = self.conv51(x)
+        x = self.relu51(x)
+        x = self.conv52(x)
+        x = self.relu52(x)
+        x = self.conv53(x)
+        x = self.relu53(x)
+        if return_tensor == 'feature5':
+            return x
+        if return_tensor == 'lpips':
+            features.append(x)
+
+        x = self.pool5(x)
+        if return_tensor == 'pool5':
+            return x
+
+        if return_tensor == 'lpips':
+            score = 0
+            assert len(features) == 5
+            for idx in range(5):
+                feature = features[idx]
+                norm = feature.norm(dim=1, keepdim=True)
+                feature = feature / (norm + 1e-10)
+                feature_x, feature_y = feature.chunk(2, dim=0)
+                diff = (feature_x - feature_y).square()
+                score += self.lpips[idx](diff).mean(dim=(2, 3), keepdim=False)
+            return score.sum(dim=1, keepdim=False)
+
+        x = self.avgpool(x)
+        x = self.flatten(x)
+        if return_tensor == 'flatten':
+            return x
+
+        x = self.fc1(x)
+        x = self.fc1_relu(x)
+        x = self.fc1_dropout(x)
+        x = self.fc2(x)
+        x = self.fc2_relu(x)
+        x = self.fc2_dropout(x)
+        if return_tensor == 'feature':
+            return x
+
+        x = self.fc3(x)
+        if return_tensor == 'logits':
+            return x
+
+        x = self.out(x)
+        if return_tensor == 'prediction':
+            return x
+
+        raise NotImplementedError(f'Output tensor name `{return_tensor}` is '
+                                  f'not implemented!')
+
+# pylint: enable=missing-function-docstring

models/pggan_discriminator.py

@@ -0,0 +1,465 @@
+# python3.7
+"""Contains the implementation of discriminator described in PGGAN.
+
+Paper: https://arxiv.org/pdf/1710.10196.pdf
+
+Official TensorFlow implementation:
+https://github.com/tkarras/progressive_growing_of_gans
+"""
+
+import numpy as np
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+__all__ = ['PGGANDiscriminator']
+
+# Resolutions allowed.
+_RESOLUTIONS_ALLOWED = [8, 16, 32, 64, 128, 256, 512, 1024]
+
+# Default gain factor for weight scaling.
+_WSCALE_GAIN = np.sqrt(2.0)
+
+# pylint: disable=missing-function-docstring
+
+class PGGANDiscriminator(nn.Module):
+    """Defines the discriminator network in PGGAN.
+
+    NOTE: The discriminator takes images with `RGB` channel order and pixel
+    range [-1, 1] as inputs.
+
+    Settings for the network:
+
+    (1) resolution: The resolution of the input image.
+    (2) init_res: Smallest resolution of the convolutional backbone.
+        (default: 4)
+    (3) image_channels: Number of channels of the input image. (default: 3)
+    (4) label_dim: Dimension of the additional label for conditional generation.
+        In one-hot conditioning case, it is equal to the number of classes. If
+        set to 0, conditioning training will be disabled. (default: 0)
+    (5) fused_scale: Whether to fused `conv2d` and `downsample` together,
+        resulting in `conv2d` with strides. (default: False)
+    (6) use_wscale: Whether to use weight scaling. (default: True)
+    (7) wscale_gain: The factor to control weight scaling. (default: sqrt(2.0))
+    (8) mbstd_groups: Group size for the minibatch standard deviation layer.
+        `0` means disable. (default: 16)
+    (9) fmaps_base: Factor to control number of feature maps for each layer.
+        (default: 16 << 10)
+    (10) fmaps_max: Maximum number of feature maps in each layer. (default: 512)
+    (11) eps: A small value to avoid divide overflow. (default: 1e-8)
+    """
+
+    def __init__(self,
+                 resolution,
+                 init_res=4,
+                 image_channels=3,
+                 label_dim=0,
+                 fused_scale=False,
+                 use_wscale=True,
+                 wscale_gain=np.sqrt(2.0),
+                 mbstd_groups=16,
+                 fmaps_base=16 << 10,
+                 fmaps_max=512,
+                 eps=1e-8):
+        """Initializes with basic settings.
+
+        Raises:
+            ValueError: If the `resolution` is not supported.
+        """
+        super().__init__()
+
+        if resolution not in _RESOLUTIONS_ALLOWED:
+            raise ValueError(f'Invalid resolution: `{resolution}`!\n'
+                             f'Resolutions allowed: {_RESOLUTIONS_ALLOWED}.')
+
+        self.init_res = init_res
+        self.init_res_log2 = int(np.log2(self.init_res))
+        self.resolution = resolution
+        self.final_res_log2 = int(np.log2(self.resolution))
+        self.image_channels = image_channels
+        self.label_dim = label_dim
+        self.fused_scale = fused_scale
+        self.use_wscale = use_wscale
+        self.wscale_gain = wscale_gain
+        self.mbstd_groups = mbstd_groups
+        self.fmaps_base = fmaps_base
+        self.fmaps_max = fmaps_max
+        self.eps = eps
+
+        # Level-of-details (used for progressive training).
+        self.register_buffer('lod', torch.zeros(()))
+        self.pth_to_tf_var_mapping = {'lod': 'lod'}
+
+        for res_log2 in range(self.final_res_log2, self.init_res_log2 - 1, -1):
+            res = 2 ** res_log2
+            in_channels = self.get_nf(res)
+            out_channels = self.get_nf(res // 2)
+            block_idx = self.final_res_log2 - res_log2
+
+            # Input convolution layer for each resolution.
+            self.add_module(
+                f'input{block_idx}',
+                ConvLayer(in_channels=self.image_channels,
+                          out_channels=in_channels,
+                          kernel_size=1,
+                          add_bias=True,
+                          downsample=False,
+                          fused_scale=False,
+                          use_wscale=use_wscale,
+                          wscale_gain=wscale_gain,
+                          activation_type='lrelu'))
+            self.pth_to_tf_var_mapping[f'input{block_idx}.weight'] = (
+                f'FromRGB_lod{block_idx}/weight')
+            self.pth_to_tf_var_mapping[f'input{block_idx}.bias'] = (
+                f'FromRGB_lod{block_idx}/bias')
+
+            # Convolution block for each resolution (except the last one).
+            if res != self.init_res:
+                self.add_module(
+                    f'layer{2 * block_idx}',
+                    ConvLayer(in_channels=in_channels,
+                              out_channels=in_channels,
+                              kernel_size=3,
+                              add_bias=True,
+                              downsample=False,
+                              fused_scale=False,
+                              use_wscale=use_wscale,
+                              wscale_gain=wscale_gain,
+                              activation_type='lrelu'))
+                tf_layer0_name = 'Conv0'
+                self.add_module(
+                    f'layer{2 * block_idx + 1}',
+                    ConvLayer(in_channels=in_channels,
+                              out_channels=out_channels,
+                              kernel_size=3,
+                              add_bias=True,
+                              downsample=True,
+                              fused_scale=fused_scale,
+                              use_wscale=use_wscale,
+                              wscale_gain=wscale_gain,
+                              activation_type='lrelu'))
+                tf_layer1_name = 'Conv1_down' if fused_scale else 'Conv1'
+
+            # Convolution block for last resolution.
+            else:
+                self.mbstd = MiniBatchSTDLayer(groups=mbstd_groups, eps=eps)
+                self.add_module(
+                    f'layer{2 * block_idx}',
+                    ConvLayer(
+                        in_channels=in_channels + 1,
+                        out_channels=in_channels,
+                        kernel_size=3,
+                        add_bias=True,
+                        downsample=False,
+                        fused_scale=False,
+                        use_wscale=use_wscale,
+                        wscale_gain=wscale_gain,
+                        activation_type='lrelu'))
+                tf_layer0_name = 'Conv'
+                self.add_module(
+                    f'layer{2 * block_idx + 1}',
+                    DenseLayer(in_channels=in_channels * res * res,
+                               out_channels=out_channels,
+                               add_bias=True,
+                               use_wscale=use_wscale,
+                               wscale_gain=wscale_gain,
+                               activation_type='lrelu'))
+                tf_layer1_name = 'Dense0'
+
+            self.pth_to_tf_var_mapping[f'layer{2 * block_idx}.weight'] = (
+                f'{res}x{res}/{tf_layer0_name}/weight')
+            self.pth_to_tf_var_mapping[f'layer{2 * block_idx}.bias'] = (
+                f'{res}x{res}/{tf_layer0_name}/bias')
+            self.pth_to_tf_var_mapping[f'layer{2 * block_idx + 1}.weight'] = (
+                f'{res}x{res}/{tf_layer1_name}/weight')
+            self.pth_to_tf_var_mapping[f'layer{2 * block_idx + 1}.bias'] = (
+                f'{res}x{res}/{tf_layer1_name}/bias')
+
+        # Final dense layer.
+        self.output = DenseLayer(in_channels=out_channels,
+                                 out_channels=1 + self.label_dim,
+                                 add_bias=True,
+                                 use_wscale=self.use_wscale,
+                                 wscale_gain=1.0,
+                                 activation_type='linear')
+        self.pth_to_tf_var_mapping['output.weight'] = (
+            f'{res}x{res}/Dense1/weight')
+        self.pth_to_tf_var_mapping['output.bias'] = (
+            f'{res}x{res}/Dense1/bias')
+
+    def get_nf(self, res):
+        """Gets number of feature maps according to the given resolution."""
+        return min(self.fmaps_base // res, self.fmaps_max)
+
+    def forward(self, image, lod=None):
+        expected_shape = (self.image_channels, self.resolution, self.resolution)
+        if image.ndim != 4 or image.shape[1:] != expected_shape:
+            raise ValueError(f'The input tensor should be with shape '
+                             f'[batch_size, channel, height, width], where '
+                             f'`channel` equals to {self.image_channels}, '
+                             f'`height`, `width` equal to {self.resolution}!\n'
+                             f'But `{image.shape}` is received!')
+
+        lod = self.lod.item() if lod is None else lod
+        if lod + self.init_res_log2 > self.final_res_log2:
+            raise ValueError(f'Maximum level-of-details (lod) is '
+                             f'{self.final_res_log2 - self.init_res_log2}, '
+                             f'but `{lod}` is received!')
+
+        lod = self.lod.item()
+        for res_log2 in range(self.final_res_log2, self.init_res_log2 - 1, -1):
+            block_idx = current_lod = self.final_res_log2 - res_log2
+            if current_lod <= lod < current_lod + 1:
+                x = getattr(self, f'input{block_idx}')(image)
+            elif current_lod - 1 < lod < current_lod:
+                alpha = lod - np.floor(lod)
+                y = getattr(self, f'input{block_idx}')(image)
+                x = y * alpha + x * (1 - alpha)
+            if lod < current_lod + 1:
+                if res_log2 == self.init_res_log2:
+                    x = self.mbstd(x)
+                x = getattr(self, f'layer{2 * block_idx}')(x)
+                x = getattr(self, f'layer{2 * block_idx + 1}')(x)
+            if lod > current_lod:
+                image = F.avg_pool2d(
+                    image, kernel_size=2, stride=2, padding=0)
+        x = self.output(x)
+
+        return {'score': x}
+
+
+class MiniBatchSTDLayer(nn.Module):
+    """Implements the minibatch standard deviation layer."""
+
+    def __init__(self, groups, eps):
+        super().__init__()
+        self.groups = groups
+        self.eps = eps
+
+    def extra_repr(self):
+        return f'groups={self.groups}, epsilon={self.eps}'
+
+    def forward(self, x):
+        if self.groups <= 1:
+            return x
+
+        N, C, H, W = x.shape
+        G = min(self.groups, N)  # Number of groups.
+
+        y = x.reshape(G, -1, C, H, W)            # [GnCHW]
+        y = y - y.mean(dim=0)                    # [GnCHW]
+        y = y.square().mean(dim=0)               # [nCHW]
+        y = (y + self.eps).sqrt()                # [nCHW]
+        y = y.mean(dim=(1, 2, 3), keepdim=True)  # [n111]
+        y = y.repeat(G, 1, H, W)                 # [N1HW]
+        x = torch.cat([x, y], dim=1)             # [N(C+1)HW]
+
+        return x
+
+
+class DownsamplingLayer(nn.Module):
+    """Implements the downsampling layer.
+
+    Basically, this layer can be used to downsample feature maps with average
+    pooling.
+    """
+
+    def __init__(self, scale_factor):
+        super().__init__()
+        self.scale_factor = scale_factor
+
+    def extra_repr(self):
+        return f'factor={self.scale_factor}'
+
+    def forward(self, x):
+        if self.scale_factor <= 1:
+            return x
+        return F.avg_pool2d(x,
+                            kernel_size=self.scale_factor,
+                            stride=self.scale_factor,
+                            padding=0)
+
+
+class ConvLayer(nn.Module):
+    """Implements the convolutional layer.
+
+    Basically, this layer executes convolution, activation, and downsampling (if
+    needed) in sequence.
+    """
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 add_bias,
+                 downsample,
+                 fused_scale,
+                 use_wscale,
+                 wscale_gain,
+                 activation_type):
+        """Initializes with layer settings.
+
+        Args:
+            in_channels: Number of channels of the input tensor.
+            out_channels: Number of channels of the output tensor.
+            kernel_size: Size of the convolutional kernels.
+            add_bias: Whether to add bias onto the convolutional result.
+            downsample: Whether to downsample the result after convolution.
+            fused_scale: Whether to fused `conv2d` and `downsample` together,
+                resulting in `conv2d` with strides.
+            use_wscale: Whether to use weight scaling.
+            wscale_gain: Gain factor for weight scaling.
+            activation_type: Type of activation.
+        """
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        self.add_bias = add_bias
+        self.downsample = downsample
+        self.fused_scale = fused_scale
+        self.use_wscale = use_wscale
+        self.wscale_gain = wscale_gain
+        self.activation_type = activation_type
+
+        if downsample and not fused_scale:
+            self.down = DownsamplingLayer(scale_factor=2)
+        else:
+            self.down = nn.Identity()
+
+        if downsample and fused_scale:
+            self.use_stride = True
+            self.stride = 2
+            self.padding = 1
+        else:
+            self.use_stride = False
+            self.stride = 1
+            self.padding = kernel_size // 2
+
+        weight_shape = (out_channels, in_channels, kernel_size, kernel_size)
+        fan_in = kernel_size * kernel_size * in_channels
+        wscale = wscale_gain / np.sqrt(fan_in)
+        if use_wscale:
+            self.weight = nn.Parameter(torch.randn(*weight_shape))
+            self.wscale = wscale
+        else:
+            self.weight = nn.Parameter(torch.randn(*weight_shape) * wscale)
+            self.wscale = 1.0
+
+        if add_bias:
+            self.bias = nn.Parameter(torch.zeros(out_channels))
+        else:
+            self.bias = None
+
+        assert activation_type in ['linear', 'relu', 'lrelu']
+
+    def extra_repr(self):
+        return (f'in_ch={self.in_channels}, '
+                f'out_ch={self.out_channels}, '
+                f'ksize={self.kernel_size}, '
+                f'wscale_gain={self.wscale_gain:.3f}, '
+                f'bias={self.add_bias}, '
+                f'downsample={self.scale_factor}, '
+                f'fused_scale={self.fused_scale}, '
+                f'act={self.activation_type}')
+
+    def forward(self, x):
+        weight = self.weight
+        if self.wscale != 1.0:
+            weight = weight * self.wscale
+
+        if self.use_stride:
+            weight = F.pad(weight, (1, 1, 1, 1, 0, 0, 0, 0), 'constant', 0.0)
+            weight = (weight[:, :, 1:, 1:] + weight[:, :, :-1, 1:] +
+                      weight[:, :, 1:, :-1] + weight[:, :, :-1, :-1]) * 0.25
+        x = F.conv2d(x,
+                     weight=weight,
+                     bias=self.bias,
+                     stride=self.stride,
+                     padding=self.padding)
+
+        if self.activation_type == 'linear':
+            pass
+        elif self.activation_type == 'relu':
+            x = F.relu(x, inplace=True)
+        elif self.activation_type == 'lrelu':
+            x = F.leaky_relu(x, negative_slope=0.2, inplace=True)
+        else:
+            raise NotImplementedError(f'Not implemented activation type '
+                                      f'`{self.activation_type}`!')
+        x = self.down(x)
+
+        return x
+
+
+class DenseLayer(nn.Module):
+    """Implements the dense layer."""
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 add_bias,
+                 use_wscale,
+                 wscale_gain,
+                 activation_type):
+        """Initializes with layer settings.
+
+        Args:
+            in_channels: Number of channels of the input tensor.
+            out_channels: Number of channels of the output tensor.
+            add_bias: Whether to add bias onto the fully-connected result.
+            use_wscale: Whether to use weight scaling.
+            wscale_gain: Gain factor for weight scaling.
+            activation_type: Type of activation.
+
+        Raises:
+            NotImplementedError: If the `activation_type` is not supported.
+        """
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.add_bias = add_bias
+        self.use_wscale = use_wscale
+        self.wscale_gain = wscale_gain
+        self.activation_type = activation_type
+
+        weight_shape = (out_channels, in_channels)
+        wscale = wscale_gain / np.sqrt(in_channels)
+        if use_wscale:
+            self.weight = nn.Parameter(torch.randn(*weight_shape))
+            self.wscale = wscale
+        else:
+            self.weight = nn.Parameter(torch.randn(*weight_shape) * wscale)
+            self.wscale = 1.0
+
+        if add_bias:
+            self.bias = nn.Parameter(torch.zeros(out_channels))
+        else:
+            self.bias = None
+
+        assert activation_type in ['linear', 'relu', 'lrelu']
+
+    def forward(self, x):
+        if x.ndim != 2:
+            x = x.flatten(start_dim=1)
+
+        weight = self.weight
+        if self.wscale != 1.0:
+            weight = weight * self.wscale
+
+        x = F.linear(x, weight=weight, bias=self.bias)
+
+        if self.activation_type == 'linear':
+            pass
+        elif self.activation_type == 'relu':
+            x = F.relu(x, inplace=True)
+        elif self.activation_type == 'lrelu':
+            x = F.leaky_relu(x, negative_slope=0.2, inplace=True)
+        else:
+            raise NotImplementedError(f'Not implemented activation type '
+                                      f'`{self.activation_type}`!')
+
+        return x
+
+# pylint: enable=missing-function-docstring

models/pggan_generator.py

@@ -0,0 +1,401 @@
+# python3.7
+"""Contains the implementation of generator described in PGGAN.
+
+Paper: https://arxiv.org/pdf/1710.10196.pdf
+
+Official TensorFlow implementation:
+https://github.com/tkarras/progressive_growing_of_gans
+"""
+
+import numpy as np
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+__all__ = ['PGGANGenerator']
+
+# Resolutions allowed.
+_RESOLUTIONS_ALLOWED = [8, 16, 32, 64, 128, 256, 512, 1024]
+
+# pylint: disable=missing-function-docstring
+
+class PGGANGenerator(nn.Module):
+    """Defines the generator network in PGGAN.
+
+    NOTE: The synthesized images are with `RGB` channel order and pixel range
+    [-1, 1].
+
+    Settings for the network:
+
+    (1) resolution: The resolution of the output image.
+    (2) init_res: The initial resolution to start with convolution. (default: 4)
+    (3) z_dim: Dimension of the input latent space, Z. (default: 512)
+    (4) image_channels: Number of channels of the output image. (default: 3)
+    (5) final_tanh: Whether to use `tanh` to control the final pixel range.
+        (default: False)
+    (6) label_dim: Dimension of the additional label for conditional generation.
+        In one-hot conditioning case, it is equal to the number of classes. If
+        set to 0, conditioning training will be disabled. (default: 0)
+    (7) fused_scale: Whether to fused `upsample` and `conv2d` together,
+        resulting in `conv2d_transpose`. (default: False)
+    (8) use_wscale: Whether to use weight scaling. (default: True)
+    (9) wscale_gain: The factor to control weight scaling. (default: sqrt(2.0))
+    (10) fmaps_base: Factor to control number of feature maps for each layer.
+         (default: 16 << 10)
+    (11) fmaps_max: Maximum number of feature maps in each layer. (default: 512)
+    (12) eps: A small value to avoid divide overflow. (default: 1e-8)
+    """
+
+    def __init__(self,
+                 resolution,
+                 init_res=4,
+                 z_dim=512,
+                 image_channels=3,
+                 final_tanh=False,
+                 label_dim=0,
+                 fused_scale=False,
+                 use_wscale=True,
+                 wscale_gain=np.sqrt(2.0),
+                 fmaps_base=16 << 10,
+                 fmaps_max=512,
+                 eps=1e-8):
+        """Initializes with basic settings.
+
+        Raises:
+            ValueError: If the `resolution` is not supported.
+        """
+        super().__init__()
+
+        if resolution not in _RESOLUTIONS_ALLOWED:
+            raise ValueError(f'Invalid resolution: `{resolution}`!\n'
+                             f'Resolutions allowed: {_RESOLUTIONS_ALLOWED}.')
+
+        self.init_res = init_res
+        self.init_res_log2 = int(np.log2(self.init_res))
+        self.resolution = resolution
+        self.final_res_log2 = int(np.log2(self.resolution))
+        self.z_dim = z_dim
+        self.image_channels = image_channels
+        self.final_tanh = final_tanh
+        self.label_dim = label_dim
+        self.fused_scale = fused_scale
+        self.use_wscale = use_wscale
+        self.wscale_gain = wscale_gain
+        self.fmaps_base = fmaps_base
+        self.fmaps_max = fmaps_max
+        self.eps = eps
+
+        # Dimension of latent space, which is convenient for sampling.
+        self.latent_dim = (self.z_dim,)
+
+        # Number of convolutional layers.
+        self.num_layers = (self.final_res_log2 - self.init_res_log2 + 1) * 2
+
+        # Level-of-details (used for progressive training).
+        self.register_buffer('lod', torch.zeros(()))
+        self.pth_to_tf_var_mapping = {'lod': 'lod'}
+
+        for res_log2 in range(self.init_res_log2, self.final_res_log2 + 1):
+            res = 2 ** res_log2
+            in_channels = self.get_nf(res // 2)
+            out_channels = self.get_nf(res)
+            block_idx = res_log2 - self.init_res_log2
+
+            # First convolution layer for each resolution.
+            if res == self.init_res:
+                self.add_module(
+                    f'layer{2 * block_idx}',
+                    ConvLayer(in_channels=z_dim + label_dim,
+                              out_channels=out_channels,
+                              kernel_size=init_res,
+                              padding=init_res - 1,
+                              add_bias=True,
+                              upsample=False,
+                              fused_scale=False,
+                              use_wscale=use_wscale,
+                              wscale_gain=wscale_gain,
+                              activation_type='lrelu',
+                              eps=eps))
+                tf_layer_name = 'Dense'
+            else:
+                self.add_module(
+                    f'layer{2 * block_idx}',
+                    ConvLayer(in_channels=in_channels,
+                              out_channels=out_channels,
+                              kernel_size=3,
+                              padding=1,
+                              add_bias=True,
+                              upsample=True,
+                              fused_scale=fused_scale,
+                              use_wscale=use_wscale,
+                              wscale_gain=wscale_gain,
+                              activation_type='lrelu',
+                              eps=eps))
+                tf_layer_name = 'Conv0_up' if fused_scale else 'Conv0'
+            self.pth_to_tf_var_mapping[f'layer{2 * block_idx}.weight'] = (
+                f'{res}x{res}/{tf_layer_name}/weight')
+            self.pth_to_tf_var_mapping[f'layer{2 * block_idx}.bias'] = (
+                f'{res}x{res}/{tf_layer_name}/bias')
+
+            # Second convolution layer for each resolution.
+            self.add_module(
+                f'layer{2 * block_idx + 1}',
+                ConvLayer(in_channels=out_channels,
+                          out_channels=out_channels,
+                          kernel_size=3,
+                          padding=1,
+                          add_bias=True,
+                          upsample=False,
+                          fused_scale=False,
+                          use_wscale=use_wscale,
+                          wscale_gain=wscale_gain,
+                          activation_type='lrelu',
+                          eps=eps))
+            tf_layer_name = 'Conv' if res == self.init_res else 'Conv1'
+            self.pth_to_tf_var_mapping[f'layer{2 * block_idx + 1}.weight'] = (
+                f'{res}x{res}/{tf_layer_name}/weight')
+            self.pth_to_tf_var_mapping[f'layer{2 * block_idx + 1}.bias'] = (
+                f'{res}x{res}/{tf_layer_name}/bias')
+
+            # Output convolution layer for each resolution.
+            self.add_module(
+                f'output{block_idx}',
+                ConvLayer(in_channels=out_channels,
+                          out_channels=image_channels,
+                          kernel_size=1,
+                          padding=0,
+                          add_bias=True,
+                          upsample=False,
+                          fused_scale=False,
+                          use_wscale=use_wscale,
+                          wscale_gain=1.0,
+                          activation_type='linear',
+                          eps=eps))
+            self.pth_to_tf_var_mapping[f'output{block_idx}.weight'] = (
+                f'ToRGB_lod{self.final_res_log2 - res_log2}/weight')
+            self.pth_to_tf_var_mapping[f'output{block_idx}.bias'] = (
+                f'ToRGB_lod{self.final_res_log2 - res_log2}/bias')
+
+    def get_nf(self, res):
+        """Gets number of feature maps according to the given resolution."""
+        return min(self.fmaps_base // res, self.fmaps_max)
+
+    def forward(self, z, label=None, lod=None):
+        if z.ndim != 2 or z.shape[1] != self.z_dim:
+            raise ValueError(f'Input latent code should be with shape '
+                             f'[batch_size, latent_dim], where '
+                             f'`latent_dim` equals to {self.z_dim}!\n'
+                             f'But `{z.shape}` is received!')
+        z = self.layer0.pixel_norm(z)
+        if self.label_dim:
+            if label is None:
+                raise ValueError(f'Model requires an additional label '
+                                 f'(with size {self.label_dim}) as input, '
+                                 f'but no label is received!')
+            if label.ndim != 2 or label.shape != (z.shape[0], self.label_dim):
+                raise ValueError(f'Input label should be with shape '
+                                 f'[batch_size, label_dim], where '
+                                 f'`batch_size` equals to that of '
+                                 f'latent codes ({z.shape[0]}) and '
+                                 f'`label_dim` equals to {self.label_dim}!\n'
+                                 f'But `{label.shape}` is received!')
+            label = label.to(dtype=torch.float32)
+            z = torch.cat((z, label), dim=1)
+
+        lod = self.lod.item() if lod is None else lod
+        if lod + self.init_res_log2 > self.final_res_log2:
+            raise ValueError(f'Maximum level-of-details (lod) is '
+                             f'{self.final_res_log2 - self.init_res_log2}, '
+                             f'but `{lod}` is received!')
+
+        x = z.view(z.shape[0], self.z_dim + self.label_dim, 1, 1)
+        for res_log2 in range(self.init_res_log2, self.final_res_log2 + 1):
+            current_lod = self.final_res_log2 - res_log2
+            block_idx = res_log2 - self.init_res_log2
+            if lod < current_lod + 1:
+                x = getattr(self, f'layer{2 * block_idx}')(x)
+                x = getattr(self, f'layer{2 * block_idx + 1}')(x)
+            if current_lod - 1 < lod <= current_lod:
+                image = getattr(self, f'output{block_idx}')(x)
+            elif current_lod < lod < current_lod + 1:
+                alpha = np.ceil(lod) - lod
+                temp = getattr(self, f'output{block_idx}')(x)
+                image = F.interpolate(image, scale_factor=2, mode='nearest')
+                image = temp * alpha + image * (1 - alpha)
+            elif lod >= current_lod + 1:
+                image = F.interpolate(image, scale_factor=2, mode='nearest')
+        if self.final_tanh:
+            image = torch.tanh(image)
+
+        results = {
+            'z': z,
+            'label': label,
+            'image': image,
+        }
+        return results
+
+
+class PixelNormLayer(nn.Module):
+    """Implements pixel-wise feature vector normalization layer."""
+
+    def __init__(self, dim, eps):
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+
+    def extra_repr(self):
+        return f'dim={self.dim}, epsilon={self.eps}'
+
+    def forward(self, x):
+        scale = (x.square().mean(dim=self.dim, keepdim=True) + self.eps).rsqrt()
+        return x * scale
+
+
+class UpsamplingLayer(nn.Module):
+    """Implements the upsampling layer.
+
+    Basically, this layer can be used to upsample feature maps with nearest
+    neighbor interpolation.
+    """
+
+    def __init__(self, scale_factor):
+        super().__init__()
+        self.scale_factor = scale_factor
+
+    def extra_repr(self):
+        return f'factor={self.scale_factor}'
+
+    def forward(self, x):
+        if self.scale_factor <= 1:
+            return x
+        return F.interpolate(x, scale_factor=self.scale_factor, mode='nearest')
+
+
+class ConvLayer(nn.Module):
+    """Implements the convolutional layer.
+
+    Basically, this layer executes pixel-wise normalization, upsampling (if
+    needed), convolution, and activation in sequence.
+    """
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 padding,
+                 add_bias,
+                 upsample,
+                 fused_scale,
+                 use_wscale,
+                 wscale_gain,
+                 activation_type,
+                 eps):
+        """Initializes with layer settings.
+
+        Args:
+            in_channels: Number of channels of the input tensor.
+            out_channels: Number of channels of the output tensor.
+            kernel_size: Size of the convolutional kernels.
+            padding: Padding used in convolution.
+            add_bias: Whether to add bias onto the convolutional result.
+            upsample: Whether to upsample the input tensor before convolution.
+            fused_scale: Whether to fused `upsample` and `conv2d` together,
+                resulting in `conv2d_transpose`.
+            use_wscale: Whether to use weight scaling.
+            wscale_gain: Gain factor for weight scaling.
+            activation_type: Type of activation.
+            eps: A small value to avoid divide overflow.
+        """
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        self.padding = padding
+        self.add_bias = add_bias
+        self.upsample = upsample
+        self.fused_scale = fused_scale
+        self.use_wscale = use_wscale
+        self.wscale_gain = wscale_gain
+        self.activation_type = activation_type
+        self.eps = eps
+
+        self.pixel_norm = PixelNormLayer(dim=1, eps=eps)
+
+        if upsample and not fused_scale:
+            self.up = UpsamplingLayer(scale_factor=2)
+        else:
+            self.up = nn.Identity()
+
+        if upsample and fused_scale:
+            self.use_conv2d_transpose = True
+            weight_shape = (in_channels, out_channels, kernel_size, kernel_size)
+            self.stride = 2
+            self.padding = 1
+        else:
+            self.use_conv2d_transpose = False
+            weight_shape = (out_channels, in_channels, kernel_size, kernel_size)
+            self.stride = 1
+
+        fan_in = kernel_size * kernel_size * in_channels
+        wscale = wscale_gain / np.sqrt(fan_in)
+        if use_wscale:
+            self.weight = nn.Parameter(torch.randn(*weight_shape))
+            self.wscale = wscale
+        else:
+            self.weight = nn.Parameter(torch.randn(*weight_shape) * wscale)
+            self.wscale = 1.0
+
+        if add_bias:
+            self.bias = nn.Parameter(torch.zeros(out_channels))
+        else:
+            self.bias = None
+
+        assert activation_type in ['linear', 'relu', 'lrelu']
+
+    def extra_repr(self):
+        return (f'in_ch={self.in_channels}, '
+                f'out_ch={self.out_channels}, '
+                f'ksize={self.kernel_size}, '
+                f'padding={self.padding}, '
+                f'wscale_gain={self.wscale_gain:.3f}, '
+                f'bias={self.add_bias}, '
+                f'upsample={self.scale_factor}, '
+                f'fused_scale={self.fused_scale}, '
+                f'act={self.activation_type}')
+
+    def forward(self, x):
+        x = self.pixel_norm(x)
+        x = self.up(x)
+        weight = self.weight
+        if self.wscale != 1.0:
+            weight = weight * self.wscale
+        if self.use_conv2d_transpose:
+            weight = F.pad(weight, (1, 1, 1, 1, 0, 0, 0, 0), 'constant', 0.0)
+            weight = (weight[:, :, 1:, 1:] + weight[:, :, :-1, 1:] +
+                      weight[:, :, 1:, :-1] + weight[:, :, :-1, :-1])
+            x = F.conv_transpose2d(x,
+                                   weight=weight,
+                                   bias=self.bias,
+                                   stride=self.stride,
+                                   padding=self.padding)
+        else:
+            x = F.conv2d(x,
+                         weight=weight,
+                         bias=self.bias,
+                         stride=self.stride,
+                         padding=self.padding)
+
+        if self.activation_type == 'linear':
+            pass
+        elif self.activation_type == 'relu':
+            x = F.relu(x, inplace=True)
+        elif self.activation_type == 'lrelu':
+            x = F.leaky_relu(x, negative_slope=0.2, inplace=True)
+        else:
+            raise NotImplementedError(f'Not implemented activation type '
+                                      f'`{self.activation_type}`!')
+
+        return x
+
+# pylint: enable=missing-function-docstring

models/stylegan2_discriminator.py

@@ -0,0 +1,729 @@
+# python3.7
+"""Contains the implementation of discriminator described in StyleGAN2.
+
+Compared to that of StyleGAN, the discriminator in StyleGAN2 mainly adds skip
+connections, increases model size and disables progressive growth. This script
+ONLY supports config F in the original paper.
+
+Paper: https://arxiv.org/pdf/1912.04958.pdf
+
+Official TensorFlow implementation: https://github.com/NVlabs/stylegan2
+"""
+
+import numpy as np
+
+import torch
+import torch.nn as nn
+
+from third_party.stylegan2_official_ops import bias_act
+from third_party.stylegan2_official_ops import upfirdn2d
+from third_party.stylegan2_official_ops import conv2d_gradfix
+
+__all__ = ['StyleGAN2Discriminator']
+
+# Resolutions allowed.
+_RESOLUTIONS_ALLOWED = [8, 16, 32, 64, 128, 256, 512, 1024]
+
+# Architectures allowed.
+_ARCHITECTURES_ALLOWED = ['resnet', 'skip', 'origin']
+
+# pylint: disable=missing-function-docstring
+
+class StyleGAN2Discriminator(nn.Module):
+    """Defines the discriminator network in StyleGAN2.
+
+    NOTE: The discriminator takes images with `RGB` channel order and pixel
+    range [-1, 1] as inputs.
+
+    Settings for the backbone:
+
+    (1) resolution: The resolution of the input image. (default: -1)
+    (2) init_res: Smallest resolution of the convolutional backbone.
+        (default: 4)
+    (3) image_channels: Number of channels of the input image. (default: 3)
+    (4) architecture: Type of architecture. Support `origin`, `skip`, and
+        `resnet`. (default: `resnet`)
+    (5) use_wscale: Whether to use weight scaling. (default: True)
+    (6) wscale_gain: The factor to control weight scaling. (default: 1.0)
+    (7) lr_mul: Learning rate multiplier for backbone. (default: 1.0)
+    (8) mbstd_groups: Group size for the minibatch standard deviation layer.
+        `0` means disable. (default: 4)
+    (9) mbstd_channels: Number of new channels (appended to the original feature
+        map) after the minibatch standard deviation layer. (default: 1)
+    (10) fmaps_base: Factor to control number of feature maps for each layer.
+         (default: 32 << 10)
+    (11) fmaps_max: Maximum number of feature maps in each layer. (default: 512)
+    (12) filter_kernel: Kernel used for filtering (e.g., downsampling).
+         (default: (1, 3, 3, 1))
+    (13) conv_clamp: A threshold to clamp the output of convolution layers to
+         avoid overflow under FP16 training. (default: None)
+    (14) eps: A small value to avoid divide overflow. (default: 1e-8)
+
+    Settings for conditional model:
+
+    (1) label_dim: Dimension of the additional label for conditional generation.
+        In one-hot conditioning case, it is equal to the number of classes. If
+        set to 0, conditioning training will be disabled. (default: 0)
+    (2) embedding_dim: Dimension of the embedding space, if needed.
+        (default: 512)
+    (3) embedding_bias: Whether to add bias to embedding learning.
+        (default: True)
+    (4) embedding_use_wscale: Whether to use weight scaling for embedding
+        learning. (default: True)
+    (5) embedding_lr_mul: Learning rate multiplier for the embedding learning.
+        (default: 1.0)
+    (6) normalize_embedding: Whether to normalize the embedding. (default: True)
+    (7) mapping_layers: Number of layers of the additional mapping network after
+        embedding. (default: 0)
+    (8) mapping_fmaps: Number of hidden channels of the additional mapping
+        network after embedding. (default: 512)
+    (9) mapping_use_wscale: Whether to use weight scaling for the additional
+        mapping network. (default: True)
+    (10) mapping_lr_mul: Learning rate multiplier for the additional mapping
+         network after embedding. (default: 0.1)
+
+    Runtime settings:
+
+    (1) fp16_res: Layers at resolution higher than (or equal to) this field will
+        use `float16` precision for computation. This is merely used for
+        acceleration. If set as `None`, all layers will use `float32` by
+        default. (default: None)
+    (2) impl: Implementation mode of some particular ops, e.g., `filtering`,
+        `bias_act`, etc. `cuda` means using the official CUDA implementation
+        from StyleGAN2, while `ref` means using the native PyTorch ops.
+        (default: `cuda`)
+    """
+
+    def __init__(self,
+                 # Settings for backbone.
+                 resolution=-1,
+                 init_res=4,
+                 image_channels=3,
+                 architecture='resnet',
+                 use_wscale=True,
+                 wscale_gain=1.0,
+                 lr_mul=1.0,
+                 mbstd_groups=4,
+                 mbstd_channels=1,
+                 fmaps_base=32 << 10,
+                 fmaps_max=512,
+                 filter_kernel=(1, 3, 3, 1),
+                 conv_clamp=None,
+                 eps=1e-8,
+                 # Settings for conditional model.
+                 label_dim=0,
+                 embedding_dim=512,
+                 embedding_bias=True,
+                 embedding_use_wscale=True,
+                 embedding_lr_mul=1.0,
+                 normalize_embedding=True,
+                 mapping_layers=0,
+                 mapping_fmaps=512,
+                 mapping_use_wscale=True,
+                 mapping_lr_mul=0.1):
+        """Initializes with basic settings.
+
+        Raises:
+            ValueError: If the `resolution` is not supported, or `architecture`
+                is not supported.
+        """
+        super().__init__()
+
+        if resolution not in _RESOLUTIONS_ALLOWED:
+            raise ValueError(f'Invalid resolution: `{resolution}`!\n'
+                             f'Resolutions allowed: {_RESOLUTIONS_ALLOWED}.')
+        architecture = architecture.lower()
+        if architecture not in _ARCHITECTURES_ALLOWED:
+            raise ValueError(f'Invalid architecture: `{architecture}`!\n'
+                             f'Architectures allowed: '
+                             f'{_ARCHITECTURES_ALLOWED}.')
+
+        self.init_res = init_res
+        self.init_res_log2 = int(np.log2(init_res))
+        self.resolution = resolution
+        self.final_res_log2 = int(np.log2(resolution))
+        self.image_channels = image_channels
+        self.architecture = architecture
+        self.use_wscale = use_wscale
+        self.wscale_gain = wscale_gain
+        self.lr_mul = lr_mul
+        self.mbstd_groups = mbstd_groups
+        self.mbstd_channels = mbstd_channels
+        self.fmaps_base = fmaps_base
+        self.fmaps_max = fmaps_max
+        self.filter_kernel = filter_kernel
+        self.conv_clamp = conv_clamp
+        self.eps = eps
+
+        self.label_dim = label_dim
+        self.embedding_dim = embedding_dim
+        self.embedding_bias = embedding_bias
+        self.embedding_use_wscale = embedding_use_wscale
+        self.embedding_lr_mul = embedding_lr_mul
+        self.normalize_embedding = normalize_embedding
+        self.mapping_layers = mapping_layers
+        self.mapping_fmaps = mapping_fmaps
+        self.mapping_use_wscale = mapping_use_wscale
+        self.mapping_lr_mul = mapping_lr_mul
+
+        self.pth_to_tf_var_mapping = {}
+
+        # Embedding for conditional discrimination.
+        self.use_embedding = label_dim > 0 and embedding_dim > 0
+        if self.use_embedding:
+            self.embedding = DenseLayer(in_channels=label_dim,
+                                        out_channels=embedding_dim,
+                                        add_bias=embedding_bias,
+                                        init_bias=0.0,
+                                        use_wscale=embedding_use_wscale,
+                                        wscale_gain=wscale_gain,
+                                        lr_mul=embedding_lr_mul,
+                                        activation_type='linear')
+            self.pth_to_tf_var_mapping['embedding.weight'] = 'LabelEmbed/weight'
+            if self.embedding_bias:
+                self.pth_to_tf_var_mapping['embedding.bias'] = 'LabelEmbed/bias'
+
+            if self.normalize_embedding:
+                self.norm = PixelNormLayer(dim=1, eps=eps)
+
+            for i in range(mapping_layers):
+                in_channels = (embedding_dim if i == 0 else mapping_fmaps)
+                out_channels = (embedding_dim if i == (mapping_layers - 1) else
+                                mapping_fmaps)
+                layer_name = f'mapping{i}'
+                self.add_module(layer_name,
+                                DenseLayer(in_channels=in_channels,
+                                           out_channels=out_channels,
+                                           add_bias=True,
+                                           init_bias=0.0,
+                                           use_wscale=mapping_use_wscale,
+                                           wscale_gain=wscale_gain,
+                                           lr_mul=mapping_lr_mul,
+                                           activation_type='lrelu'))
+                self.pth_to_tf_var_mapping[f'{layer_name}.weight'] = (
+                    f'Mapping{i}/weight')
+                self.pth_to_tf_var_mapping[f'{layer_name}.bias'] = (
+                    f'Mapping{i}/bias')
+
+        # Convolutional backbone.
+        for res_log2 in range(self.final_res_log2, self.init_res_log2 - 1, -1):
+            res = 2 ** res_log2
+            in_channels = self.get_nf(res)
+            out_channels = self.get_nf(res // 2)
+            block_idx = self.final_res_log2 - res_log2
+
+            # Input convolution layer for each resolution (if needed).
+            if res_log2 == self.final_res_log2 or self.architecture == 'skip':
+                layer_name = f'input{block_idx}'
+                self.add_module(layer_name,
+                                ConvLayer(in_channels=image_channels,
+                                          out_channels=in_channels,
+                                          kernel_size=1,
+                                          add_bias=True,
+                                          scale_factor=1,
+                                          filter_kernel=None,
+                                          use_wscale=use_wscale,
+                                          wscale_gain=wscale_gain,
+                                          lr_mul=lr_mul,
+                                          activation_type='lrelu',
+                                          conv_clamp=conv_clamp))
+                self.pth_to_tf_var_mapping[f'{layer_name}.weight'] = (
+                    f'{res}x{res}/FromRGB/weight')
+                self.pth_to_tf_var_mapping[f'{layer_name}.bias'] = (
+                    f'{res}x{res}/FromRGB/bias')
+
+            # Convolution block for each resolution (except the last one).
+            if res != self.init_res:
+                # First layer (kernel 3x3) without downsampling.
+                layer_name = f'layer{2 * block_idx}'
+                self.add_module(layer_name,
+                                ConvLayer(in_channels=in_channels,
+                                          out_channels=in_channels,
+                                          kernel_size=3,
+                                          add_bias=True,
+                                          scale_factor=1,
+                                          filter_kernel=None,
+                                          use_wscale=use_wscale,
+                                          wscale_gain=wscale_gain,
+                                          lr_mul=lr_mul,
+                                          activation_type='lrelu',
+                                          conv_clamp=conv_clamp))
+                self.pth_to_tf_var_mapping[f'{layer_name}.weight'] = (
+                    f'{res}x{res}/Conv0/weight')
+                self.pth_to_tf_var_mapping[f'{layer_name}.bias'] = (
+                    f'{res}x{res}/Conv0/bias')
+
+                # Second layer (kernel 3x3) with downsampling
+                layer_name = f'layer{2 * block_idx + 1}'
+                self.add_module(layer_name,
+                                ConvLayer(in_channels=in_channels,
+                                          out_channels=out_channels,
+                                          kernel_size=3,
+                                          add_bias=True,
+                                          scale_factor=2,
+                                          filter_kernel=filter_kernel,
+                                          use_wscale=use_wscale,
+                                          wscale_gain=wscale_gain,
+                                          lr_mul=lr_mul,
+                                          activation_type='lrelu',
+                                          conv_clamp=conv_clamp))
+                self.pth_to_tf_var_mapping[f'{layer_name}.weight'] = (
+                    f'{res}x{res}/Conv1_down/weight')
+                self.pth_to_tf_var_mapping[f'{layer_name}.bias'] = (
+                    f'{res}x{res}/Conv1_down/bias')
+
+                # Residual branch (kernel 1x1) with downsampling, without bias,
+                # with linear activation.
+                if self.architecture == 'resnet':
+                    layer_name = f'residual{block_idx}'
+                    self.add_module(layer_name,
+                                    ConvLayer(in_channels=in_channels,
+                                              out_channels=out_channels,
+                                              kernel_size=1,
+                                              add_bias=False,
+                                              scale_factor=2,
+                                              filter_kernel=filter_kernel,
+                                              use_wscale=use_wscale,
+                                              wscale_gain=wscale_gain,
+                                              lr_mul=lr_mul,
+                                              activation_type='linear',
+                                              conv_clamp=None))
+                    self.pth_to_tf_var_mapping[f'{layer_name}.weight'] = (
+                        f'{res}x{res}/Skip/weight')
+
+            # Convolution block for last resolution.
+            else:
+                self.mbstd = MiniBatchSTDLayer(
+                    groups=mbstd_groups, new_channels=mbstd_channels, eps=eps)
+
+                # First layer (kernel 3x3) without downsampling.
+                layer_name = f'layer{2 * block_idx}'
+                self.add_module(
+                    layer_name,
+                    ConvLayer(in_channels=in_channels + mbstd_channels,
+                              out_channels=in_channels,
+                              kernel_size=3,
+                              add_bias=True,
+                              scale_factor=1,
+                              filter_kernel=None,
+                              use_wscale=use_wscale,
+                              wscale_gain=wscale_gain,
+                              lr_mul=lr_mul,
+                              activation_type='lrelu',
+                              conv_clamp=conv_clamp))
+                self.pth_to_tf_var_mapping[f'{layer_name}.weight'] = (
+                    f'{res}x{res}/Conv/weight')
+                self.pth_to_tf_var_mapping[f'{layer_name}.bias'] = (
+                    f'{res}x{res}/Conv/bias')
+
+                # Second layer, as a fully-connected layer.
+                layer_name = f'layer{2 * block_idx + 1}'
+                self.add_module(layer_name,
+                                DenseLayer(in_channels=in_channels * res * res,
+                                           out_channels=in_channels,
+                                           add_bias=True,
+                                           init_bias=0.0,
+                                           use_wscale=use_wscale,
+                                           wscale_gain=wscale_gain,
+                                           lr_mul=lr_mul,
+                                           activation_type='lrelu'))
+                self.pth_to_tf_var_mapping[f'{layer_name}.weight'] = (
+                    f'{res}x{res}/Dense0/weight')
+                self.pth_to_tf_var_mapping[f'{layer_name}.bias'] = (
+                    f'{res}x{res}/Dense0/bias')
+
+                # Final dense layer to output score.
+                self.output = DenseLayer(in_channels=in_channels,
+                                         out_channels=(embedding_dim
+                                                       if self.use_embedding
+                                                       else max(label_dim, 1)),
+                                         add_bias=True,
+                                         init_bias=0.0,
+                                         use_wscale=use_wscale,
+                                         wscale_gain=wscale_gain,
+                                         lr_mul=lr_mul,
+                                         activation_type='linear')
+                self.pth_to_tf_var_mapping['output.weight'] = 'Output/weight'
+                self.pth_to_tf_var_mapping['output.bias'] = 'Output/bias'
+
+        # Used for downsampling input image for `skip` architecture.
+        if self.architecture == 'skip':
+            self.register_buffer(
+                'filter', upfirdn2d.setup_filter(filter_kernel))
+
+    def get_nf(self, res):
+        """Gets number of feature maps according to the given resolution."""
+        return min(self.fmaps_base // res, self.fmaps_max)
+
+    def forward(self, image, label=None, fp16_res=None, impl='cuda'):
+        # Check shape.
+        expected_shape = (self.image_channels, self.resolution, self.resolution)
+        if image.ndim != 4 or image.shape[1:] != expected_shape:
+            raise ValueError(f'The input tensor should be with shape '
+                             f'[batch_size, channel, height, width], where '
+                             f'`channel` equals to {self.image_channels}, '
+                             f'`height`, `width` equal to {self.resolution}!\n'
+                             f'But `{image.shape}` is received!')
+        if self.label_dim > 0:
+            if label is None:
+                raise ValueError(f'Model requires an additional label '
+                                 f'(with dimension {self.label_dim}) as input, '
+                                 f'but no label is received!')
+            batch_size = image.shape[0]
+            if label.ndim != 2 or label.shape != (batch_size, self.label_dim):
+                raise ValueError(f'Input label should be with shape '
+                                 f'[batch_size, label_dim], where '
+                                 f'`batch_size` equals to that of '
+                                 f'images ({image.shape[0]}) and '
+                                 f'`label_dim` equals to {self.label_dim}!\n'
+                                 f'But `{label.shape}` is received!')
+            label = label.to(dtype=torch.float32)
+            if self.use_embedding:
+                embed = self.embedding(label, impl=impl)
+                if self.normalize_embedding:
+                    embed = self.norm(embed)
+                for i in range(self.mapping_layers):
+                    embed = getattr(self, f'mapping{i}')(embed, impl=impl)
+
+        # Cast to `torch.float16` if needed.
+        if fp16_res is not None and self.resolution >= fp16_res:
+            image = image.to(torch.float16)
+
+        x = self.input0(image, impl=impl)
+
+        for res_log2 in range(self.final_res_log2, self.init_res_log2, -1):
+            res = 2 ** res_log2
+            # Cast to `torch.float16` if needed.
+            if fp16_res is not None and res >= fp16_res:
+                x = x.to(torch.float16)
+            else:
+                x = x.to(torch.float32)
+
+            idx = self.final_res_log2 - res_log2  # Block index
+
+            if self.architecture == 'skip' and idx > 0:
+                image = upfirdn2d.downsample2d(image, self.filter, impl=impl)
+                # Cast to `torch.float16` if needed.
+                if fp16_res is not None and res >= fp16_res:
+                    image = image.to(torch.float16)
+                else:
+                    image = image.to(torch.float32)
+                y = getattr(self, f'input{idx}')(image, impl=impl)
+                x = x + y
+
+            if self.architecture == 'resnet':
+                residual = getattr(self, f'residual{idx}')(
+                    x, runtime_gain=np.sqrt(0.5), impl=impl)
+                x = getattr(self, f'layer{2 * idx}')(x, impl=impl)
+                x = getattr(self, f'layer{2 * idx + 1}')(
+                    x, runtime_gain=np.sqrt(0.5), impl=impl)
+                x = x + residual
+            else:
+                x = getattr(self, f'layer{2 * idx}')(x, impl=impl)
+                x = getattr(self, f'layer{2 * idx + 1}')(x, impl=impl)
+
+        # Final output.
+        idx += 1
+        if fp16_res is not None:  # Always use FP32 for the last block.
+            x = x.to(torch.float32)
+        if self.architecture == 'skip':
+            image = upfirdn2d.downsample2d(image, self.filter, impl=impl)
+            if fp16_res is not None:  # Always use FP32 for the last block.
+                image = image.to(torch.float32)
+            y = getattr(self, f'input{idx}')(image, impl=impl)
+            x = x + y
+        x = self.mbstd(x)
+        x = getattr(self, f'layer{2 * idx}')(x, impl=impl)
+        x = getattr(self, f'layer{2 * idx + 1}')(x, impl=impl)
+        x = self.output(x, impl=impl)
+
+        if self.use_embedding:
+            x = (x * embed).sum(dim=1, keepdim=True)
+            x = x / np.sqrt(self.embedding_dim)
+        elif self.label_dim > 0:
+            x = (x * label).sum(dim=1, keepdim=True)
+
+        results = {
+            'score': x,
+            'label': label
+        }
+        if self.use_embedding:
+            results['embedding'] = embed
+        return results
+
+
+class PixelNormLayer(nn.Module):
+    """Implements pixel-wise feature vector normalization layer."""
+
+    def __init__(self, dim, eps):
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+
+    def extra_repr(self):
+        return f'dim={self.dim}, epsilon={self.eps}'
+
+    def forward(self, x):
+        scale = (x.square().mean(dim=self.dim, keepdim=True) + self.eps).rsqrt()
+        return x * scale
+
+
+class MiniBatchSTDLayer(nn.Module):
+    """Implements the minibatch standard deviation layer."""
+
+    def __init__(self, groups, new_channels, eps):
+        super().__init__()
+        self.groups = groups
+        self.new_channels = new_channels
+        self.eps = eps
+
+    def extra_repr(self):
+        return (f'groups={self.groups}, '
+                f'new_channels={self.new_channels}, '
+                f'epsilon={self.eps}')
+
+    def forward(self, x):
+        if self.groups <= 1 or self.new_channels < 1:
+            return x
+
+        dtype = x.dtype
+
+        N, C, H, W = x.shape
+        G = min(self.groups, N)  # Number of groups.
+        nC = self.new_channels  # Number of channel groups.
+        c = C // nC             # Channels per channel group.
+
+        y = x.reshape(G, -1, nC, c, H, W)  # [GnFcHW]
+        y = y - y.mean(dim=0)              # [GnFcHW]
+        y = y.square().mean(dim=0)         # [nFcHW]
+        y = (y + self.eps).sqrt()          # [nFcHW]
+        y = y.mean(dim=(2, 3, 4))          # [nF]
+        y = y.reshape(-1, nC, 1, 1)        # [nF11]
+        y = y.repeat(G, 1, H, W)           # [NFHW]
+        x = torch.cat((x, y), dim=1)       # [N(C+F)HW]
+
+        assert x.dtype == dtype
+        return x
+
+
+class ConvLayer(nn.Module):
+    """Implements the convolutional layer.
+
+    If downsampling is needed (i.e., `scale_factor = 2`), the feature map will
+    be filtered with `filter_kernel` first.
+    """
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 add_bias,
+                 scale_factor,
+                 filter_kernel,
+                 use_wscale,
+                 wscale_gain,
+                 lr_mul,
+                 activation_type,
+                 conv_clamp):
+        """Initializes with layer settings.
+
+        Args:
+            in_channels: Number of channels of the input tensor.
+            out_channels: Number of channels of the output tensor.
+            kernel_size: Size of the convolutional kernels.
+            add_bias: Whether to add bias onto the convolutional result.
+            scale_factor: Scale factor for downsampling. `1` means skip
+                downsampling.
+            filter_kernel: Kernel used for filtering.
+            use_wscale: Whether to use weight scaling.
+            wscale_gain: Gain factor for weight scaling.
+            lr_mul: Learning multiplier for both weight and bias.
+            activation_type: Type of activation.
+            conv_clamp: A threshold to clamp the output of convolution layers to
+                avoid overflow under FP16 training.
+        """
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        self.add_bias = add_bias
+        self.scale_factor = scale_factor
+        self.filter_kernel = filter_kernel
+        self.use_wscale = use_wscale
+        self.wscale_gain = wscale_gain
+        self.lr_mul = lr_mul
+        self.activation_type = activation_type
+        self.conv_clamp = conv_clamp
+
+        weight_shape = (out_channels, in_channels, kernel_size, kernel_size)
+        fan_in = kernel_size * kernel_size * in_channels
+        wscale = wscale_gain / np.sqrt(fan_in)
+        if use_wscale:
+            self.weight = nn.Parameter(torch.randn(*weight_shape) / lr_mul)
+            self.wscale = wscale * lr_mul
+        else:
+            self.weight = nn.Parameter(
+                torch.randn(*weight_shape) * wscale / lr_mul)
+            self.wscale = lr_mul
+
+        if add_bias:
+            self.bias = nn.Parameter(torch.zeros(out_channels))
+            self.bscale = lr_mul
+        else:
+            self.bias = None
+        self.act_gain = bias_act.activation_funcs[activation_type].def_gain
+
+        if scale_factor > 1:
+            assert filter_kernel is not None
+            self.register_buffer(
+                'filter', upfirdn2d.setup_filter(filter_kernel))
+            fh, fw = self.filter.shape
+            self.filter_padding = (
+                kernel_size // 2 + (fw - scale_factor + 1) // 2,
+                kernel_size // 2 + (fw - scale_factor) // 2,
+                kernel_size // 2 + (fh - scale_factor + 1) // 2,
+                kernel_size // 2 + (fh - scale_factor) // 2)
+
+    def extra_repr(self):
+        return (f'in_ch={self.in_channels}, '
+                f'out_ch={self.out_channels}, '
+                f'ksize={self.kernel_size}, '
+                f'wscale_gain={self.wscale_gain:.3f}, '
+                f'bias={self.add_bias}, '
+                f'lr_mul={self.lr_mul:.3f}, '
+                f'downsample={self.scale_factor}, '
+                f'downsample_filter={self.filter_kernel}, '
+                f'act={self.activation_type}, '
+                f'clamp={self.conv_clamp}')
+
+    def forward(self, x, runtime_gain=1.0, impl='cuda'):
+        dtype = x.dtype
+
+        weight = self.weight
+        if self.wscale != 1.0:
+            weight = weight * self.wscale
+        bias = None
+        if self.bias is not None:
+            bias = self.bias.to(dtype)
+            if self.bscale != 1.0:
+                bias = bias * self.bscale
+
+        if self.scale_factor == 1:  # Native convolution without downsampling.
+            padding = self.kernel_size // 2
+            x = conv2d_gradfix.conv2d(
+                x, weight.to(dtype), stride=1, padding=padding, impl=impl)
+        else:  # Convolution with downsampling.
+            down = self.scale_factor
+            f = self.filter
+            padding = self.filter_padding
+            # When kernel size = 1, use filtering function for downsampling.
+            if self.kernel_size == 1:
+                x = upfirdn2d.upfirdn2d(
+                    x, f, down=down, padding=padding, impl=impl)
+                x = conv2d_gradfix.conv2d(
+                    x, weight.to(dtype), stride=1, padding=0, impl=impl)
+            # When kernel size != 1, use stride convolution for downsampling.
+            else:
+                x = upfirdn2d.upfirdn2d(
+                    x, f, down=1, padding=padding, impl=impl)
+                x = conv2d_gradfix.conv2d(
+                    x, weight.to(dtype), stride=down, padding=0, impl=impl)
+
+        act_gain = self.act_gain * runtime_gain
+        act_clamp = None
+        if self.conv_clamp is not None:
+            act_clamp = self.conv_clamp * runtime_gain
+        x = bias_act.bias_act(x, bias,
+                              act=self.activation_type,
+                              gain=act_gain,
+                              clamp=act_clamp,
+                              impl=impl)
+
+        assert x.dtype == dtype
+        return x
+
+
+class DenseLayer(nn.Module):
+    """Implements the dense layer."""
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 add_bias,
+                 init_bias,
+                 use_wscale,
+                 wscale_gain,
+                 lr_mul,
+                 activation_type):
+        """Initializes with layer settings.
+
+        Args:
+            in_channels: Number of channels of the input tensor.
+            out_channels: Number of channels of the output tensor.
+            add_bias: Whether to add bias onto the fully-connected result.
+            init_bias: The initial bias value before training.
+            use_wscale: Whether to use weight scaling.
+            wscale_gain: Gain factor for weight scaling.
+            lr_mul: Learning multiplier for both weight and bias.
+            activation_type: Type of activation.
+        """
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.add_bias = add_bias
+        self.init_bias = init_bias
+        self.use_wscale = use_wscale
+        self.wscale_gain = wscale_gain
+        self.lr_mul = lr_mul
+        self.activation_type = activation_type
+
+        weight_shape = (out_channels, in_channels)
+        wscale = wscale_gain / np.sqrt(in_channels)
+        if use_wscale:
+            self.weight = nn.Parameter(torch.randn(*weight_shape) / lr_mul)
+            self.wscale = wscale * lr_mul
+        else:
+            self.weight = nn.Parameter(
+                torch.randn(*weight_shape) * wscale / lr_mul)
+            self.wscale = lr_mul
+
+        if add_bias:
+            init_bias = np.float32(init_bias) / lr_mul
+            self.bias = nn.Parameter(torch.full([out_channels], init_bias))
+            self.bscale = lr_mul
+        else:
+            self.bias = None
+
+    def extra_repr(self):
+        return (f'in_ch={self.in_channels}, '
+                f'out_ch={self.out_channels}, '
+                f'wscale_gain={self.wscale_gain:.3f}, '
+                f'bias={self.add_bias}, '
+                f'init_bias={self.init_bias}, '
+                f'lr_mul={self.lr_mul:.3f}, '
+                f'act={self.activation_type}')
+
+    def forward(self, x, impl='cuda'):
+        dtype = x.dtype
+
+        if x.ndim != 2:
+            x = x.flatten(start_dim=1)
+
+        weight = self.weight.to(dtype) * self.wscale
+        bias = None
+        if self.bias is not None:
+            bias = self.bias.to(dtype)
+            if self.bscale != 1.0:
+                bias = bias * self.bscale
+
+        # Fast pass for linear activation.
+        if self.activation_type == 'linear' and bias is not None:
+            x = torch.addmm(bias.unsqueeze(0), x, weight.t())
+        else:
+            x = x.matmul(weight.t())
+            x = bias_act.bias_act(x, bias, act=self.activation_type, impl=impl)
+
+        assert x.dtype == dtype
+        return x
+
+# pylint: enable=missing-function-docstring

models/stylegan2_generator.py

@@ -0,0 +1,1394 @@
+# python3.7
+"""Contains the implementation of generator described in StyleGAN2.
+
+Compared to that of StyleGAN, the generator in StyleGAN2 mainly introduces style
+demodulation, adds skip connections, increases model size, and disables
+progressive growth. This script ONLY supports config F in the original paper.
+
+Paper: https://arxiv.org/pdf/1912.04958.pdf
+
+Official TensorFlow implementation: https://github.com/NVlabs/stylegan2
+"""
+
+import numpy as np
+
+import torch
+import torch.nn as nn
+
+from third_party.stylegan2_official_ops import fma
+from third_party.stylegan2_official_ops import bias_act
+from third_party.stylegan2_official_ops import upfirdn2d
+from third_party.stylegan2_official_ops import conv2d_gradfix
+from .utils.ops import all_gather
+
+__all__ = ['StyleGAN2Generator']
+
+# Resolutions allowed.
+_RESOLUTIONS_ALLOWED = [8, 16, 32, 64, 128, 256, 512, 1024]
+
+# Architectures allowed.
+_ARCHITECTURES_ALLOWED = ['resnet', 'skip', 'origin']
+
+# pylint: disable=missing-function-docstring
+
+class StyleGAN2Generator(nn.Module):
+    """Defines the generator network in StyleGAN2.
+
+    NOTE: The synthesized images are with `RGB` channel order and pixel range
+    [-1, 1].
+
+    Settings for the mapping network:
+
+    (1) z_dim: Dimension of the input latent space, Z. (default: 512)
+    (2) w_dim: Dimension of the output latent space, W. (default: 512)
+    (3) repeat_w: Repeat w-code for different layers. (default: True)
+    (4) normalize_z: Whether to normalize the z-code. (default: True)
+    (5) mapping_layers: Number of layers of the mapping network. (default: 8)
+    (6) mapping_fmaps: Number of hidden channels of the mapping network.
+        (default: 512)
+    (7) mapping_use_wscale: Whether to use weight scaling for the mapping
+        network. (default: True)
+    (8) mapping_wscale_gain: The factor to control weight scaling for the
+        mapping network (default: 1.0)
+    (9) mapping_lr_mul: Learning rate multiplier for the mapping network.
+        (default: 0.01)
+
+    Settings for conditional generation:
+
+    (1) label_dim: Dimension of the additional label for conditional generation.
+        In one-hot conditioning case, it is equal to the number of classes. If
+        set to 0, conditioning training will be disabled. (default: 0)
+    (2) embedding_dim: Dimension of the embedding space, if needed.
+        (default: 512)
+    (3) embedding_bias: Whether to add bias to embedding learning.
+        (default: True)
+    (4) embedding_use_wscale: Whether to use weight scaling for embedding
+        learning. (default: True)
+    (5) embedding_wscale_gain: The factor to control weight scaling for
+        embedding. (default: 1.0)
+    (6) embedding_lr_mul: Learning rate multiplier for the embedding learning.
+        (default: 1.0)
+    (7) normalize_embedding: Whether to normalize the embedding. (default: True)
+    (8) normalize_embedding_latent: Whether to normalize the embedding together
+        with the latent. (default: False)
+
+    Settings for the synthesis network:
+
+    (1) resolution: The resolution of the output image. (default: -1)
+    (2) init_res: The initial resolution to start with convolution. (default: 4)
+    (3) image_channels: Number of channels of the output image. (default: 3)
+    (4) final_tanh: Whether to use `tanh` to control the final pixel range.
+        (default: False)
+    (5) const_input: Whether to use a constant in the first convolutional layer.
+        (default: True)
+    (6) architecture: Type of architecture. Support `origin`, `skip`, and
+        `resnet`. (default: `skip`)
+    (7) demodulate: Whether to perform style demodulation. (default: True)
+    (8) use_wscale: Whether to use weight scaling. (default: True)
+    (9) wscale_gain: The factor to control weight scaling. (default: 1.0)
+    (10) lr_mul: Learning rate multiplier for the synthesis network.
+         (default: 1.0)
+    (11) noise_type: Type of noise added to the convolutional results at each
+         layer. (default: `spatial`)
+    (12) fmaps_base: Factor to control number of feature maps for each layer.
+         (default: 32 << 10)
+    (13) fmaps_max: Maximum number of feature maps in each layer. (default: 512)
+    (14) filter_kernel: Kernel used for filtering (e.g., downsampling).
+         (default: (1, 3, 3, 1))
+    (15) conv_clamp: A threshold to clamp the output of convolution layers to
+         avoid overflow under FP16 training. (default: None)
+    (16) eps: A small value to avoid divide overflow. (default: 1e-8)
+
+    Runtime settings:
+
+    (1) w_moving_decay: Decay factor for updating `w_avg`, which is used for
+        training only. Set `None` to disable. (default: None)
+    (2) sync_w_avg: Synchronizing the stats of `w_avg` across replicas. If set
+        as `True`, the stats will be more accurate, yet the speed maybe a little
+        bit slower. (default: False)
+    (3) style_mixing_prob: Probability to perform style mixing as a training
+        regularization. Set `None` to disable. (default: None)
+    (4) trunc_psi: Truncation psi, set `None` to disable. (default: None)
+    (5) trunc_layers: Number of layers to perform truncation. (default: None)
+    (6) noise_mode: Mode of the layer-wise noise. Support `none`, `random`,
+        `const`. (default: `const`)
+    (7) fused_modulate: Whether to fuse `style_modulate` and `conv2d` together.
+        (default: False)
+    (8) fp16_res: Layers at resolution higher than (or equal to) this field will
+        use `float16` precision for computation. This is merely used for
+        acceleration. If set as `None`, all layers will use `float32` by
+        default. (default: None)
+    (9) impl: Implementation mode of some particular ops, e.g., `filtering`,
+        `bias_act`, etc. `cuda` means using the official CUDA implementation
+        from StyleGAN2, while `ref` means using the native PyTorch ops.
+        (default: `cuda`)
+    """
+
+    def __init__(self,
+                 # Settings for mapping network.
+                 z_dim=512,
+                 w_dim=512,
+                 repeat_w=True,
+                 normalize_z=True,
+                 mapping_layers=8,
+                 mapping_fmaps=512,
+                 mapping_use_wscale=True,
+                 mapping_wscale_gain=1.0,
+                 mapping_lr_mul=0.01,
+                 # Settings for conditional generation.
+                 label_dim=0,
+                 embedding_dim=512,
+                 embedding_bias=True,
+                 embedding_use_wscale=True,
+                 embedding_wscale_gian=1.0,
+                 embedding_lr_mul=1.0,
+                 normalize_embedding=True,
+                 normalize_embedding_latent=False,
+                 # Settings for synthesis network.
+                 resolution=-1,
+                 init_res=4,
+                 image_channels=3,
+                 final_tanh=False,
+                 const_input=True,
+                 architecture='skip',
+                 demodulate=True,
+                 use_wscale=True,
+                 wscale_gain=1.0,
+                 lr_mul=1.0,
+                 noise_type='spatial',
+                 fmaps_base=32 << 10,
+                 fmaps_max=512,
+                 filter_kernel=(1, 3, 3, 1),
+                 conv_clamp=None,
+                 eps=1e-8):
+        """Initializes with basic settings.
+
+        Raises:
+            ValueError: If the `resolution` is not supported, or `architecture`
+                is not supported.
+        """
+        super().__init__()
+
+        if resolution not in _RESOLUTIONS_ALLOWED:
+            raise ValueError(f'Invalid resolution: `{resolution}`!\n'
+                             f'Resolutions allowed: {_RESOLUTIONS_ALLOWED}.')
+        architecture = architecture.lower()
+        if architecture not in _ARCHITECTURES_ALLOWED:
+            raise ValueError(f'Invalid architecture: `{architecture}`!\n'
+                             f'Architectures allowed: '
+                             f'{_ARCHITECTURES_ALLOWED}.')
+
+        self.z_dim = z_dim
+        self.w_dim = w_dim
+        self.repeat_w = repeat_w
+        self.normalize_z = normalize_z
+        self.mapping_layers = mapping_layers
+        self.mapping_fmaps = mapping_fmaps
+        self.mapping_use_wscale = mapping_use_wscale
+        self.mapping_wscale_gain = mapping_wscale_gain
+        self.mapping_lr_mul = mapping_lr_mul
+
+        self.label_dim = label_dim
+        self.embedding_dim = embedding_dim
+        self.embedding_bias = embedding_bias
+        self.embedding_use_wscale = embedding_use_wscale
+        self.embedding_wscale_gain = embedding_wscale_gian
+        self.embedding_lr_mul = embedding_lr_mul
+        self.normalize_embedding = normalize_embedding
+        self.normalize_embedding_latent = normalize_embedding_latent
+
+        self.resolution = resolution
+        self.init_res = init_res
+        self.image_channels = image_channels
+        self.final_tanh = final_tanh
+        self.const_input = const_input
+        self.architecture = architecture
+        self.demodulate = demodulate
+        self.use_wscale = use_wscale
+        self.wscale_gain = wscale_gain
+        self.lr_mul = lr_mul
+        self.noise_type = noise_type.lower()
+        self.fmaps_base = fmaps_base
+        self.fmaps_max = fmaps_max
+        self.filter_kernel = filter_kernel
+        self.conv_clamp = conv_clamp
+        self.eps = eps
+
+        # Dimension of latent space, which is convenient for sampling.
+        self.latent_dim = (z_dim,)
+
+        # Number of synthesis (convolutional) layers.
+        self.num_layers = int(np.log2(resolution // init_res * 2)) * 2
+
+        self.mapping = MappingNetwork(
+            input_dim=z_dim,
+            output_dim=w_dim,
+            num_outputs=self.num_layers,
+            repeat_output=repeat_w,
+            normalize_input=normalize_z,
+            num_layers=mapping_layers,
+            hidden_dim=mapping_fmaps,
+            use_wscale=mapping_use_wscale,
+            wscale_gain=mapping_wscale_gain,
+            lr_mul=mapping_lr_mul,
+            label_dim=label_dim,
+            embedding_dim=embedding_dim,
+            embedding_bias=embedding_bias,
+            embedding_use_wscale=embedding_use_wscale,
+            embedding_wscale_gian=embedding_wscale_gian,
+            embedding_lr_mul=embedding_lr_mul,
+            normalize_embedding=normalize_embedding,
+            normalize_embedding_latent=normalize_embedding_latent,
+            eps=eps)
+
+        # This is used for truncation trick.
+        if self.repeat_w:
+            self.register_buffer('w_avg', torch.zeros(w_dim))
+        else:
+            self.register_buffer('w_avg', torch.zeros(self.num_layers * w_dim))
+
+        self.synthesis = SynthesisNetwork(resolution=resolution,
+                                          init_res=init_res,
+                                          w_dim=w_dim,
+                                          image_channels=image_channels,
+                                          final_tanh=final_tanh,
+                                          const_input=const_input,
+                                          architecture=architecture,
+                                          demodulate=demodulate,
+                                          use_wscale=use_wscale,
+                                          wscale_gain=wscale_gain,
+                                          lr_mul=lr_mul,
+                                          noise_type=noise_type,
+                                          fmaps_base=fmaps_base,
+                                          filter_kernel=filter_kernel,
+                                          fmaps_max=fmaps_max,
+                                          conv_clamp=conv_clamp,
+                                          eps=eps)
+
+        self.pth_to_tf_var_mapping = {'w_avg': 'dlatent_avg'}
+        for key, val in self.mapping.pth_to_tf_var_mapping.items():
+            self.pth_to_tf_var_mapping[f'mapping.{key}'] = val
+        for key, val in self.synthesis.pth_to_tf_var_mapping.items():
+            self.pth_to_tf_var_mapping[f'synthesis.{key}'] = val
+
+    def set_space_of_latent(self, space_of_latent):
+        """Sets the space to which the latent code belong.
+
+        See `SynthesisNetwork` for more details.
+        """
+        self.synthesis.set_space_of_latent(space_of_latent)
+
+    def forward(self,
+                z,
+                label=None,
+                w_moving_decay=None,
+                sync_w_avg=False,
+                style_mixing_prob=None,
+                trunc_psi=None,
+                trunc_layers=None,
+                noise_mode='const',
+                fused_modulate=False,
+                fp16_res=None,
+                impl='cuda'):
+        """Connects mapping network and synthesis network.
+
+        This forward function will also update the average `w_code`, perform
+        style mixing as a training regularizer, and do truncation trick, which
+        is specially designed for inference.
+
+        Concretely, the truncation trick acts as follows:
+
+        For layers in range [0, truncation_layers), the truncated w-code is
+        computed as
+
+        w_new = w_avg + (w - w_avg) * truncation_psi
+
+        To disable truncation, please set
+
+        (1) truncation_psi = 1.0 (None) OR
+        (2) truncation_layers = 0 (None)
+        """
+
+        mapping_results = self.mapping(z, label, impl=impl)
+
+        w = mapping_results['w']
+        if self.training and w_moving_decay is not None:
+            if sync_w_avg:
+                batch_w_avg = all_gather(w.detach()).mean(dim=0)
+            else:
+                batch_w_avg = w.detach().mean(dim=0)
+            self.w_avg.copy_(batch_w_avg.lerp(self.w_avg, w_moving_decay))
+
+        wp = mapping_results.pop('wp')
+        if self.training and style_mixing_prob is not None:
+            if np.random.uniform() < style_mixing_prob:
+                new_z = torch.randn_like(z)
+                new_wp = self.mapping(new_z, label, impl=impl)['wp']
+                mixing_cutoff = np.random.randint(1, self.num_layers)
+                wp[:, mixing_cutoff:] = new_wp[:, mixing_cutoff:]
+
+        if not self.training:
+            trunc_psi = 1.0 if trunc_psi is None else trunc_psi
+            trunc_layers = 0 if trunc_layers is None else trunc_layers
+            if trunc_psi < 1.0 and trunc_layers > 0:
+                w_avg = self.w_avg.reshape(1, -1, self.w_dim)[:, :trunc_layers]
+                wp[:, :trunc_layers] = w_avg.lerp(
+                    wp[:, :trunc_layers], trunc_psi)
+
+        synthesis_results = self.synthesis(wp,
+                                           noise_mode=noise_mode,
+                                           fused_modulate=fused_modulate,
+                                           impl=impl,
+                                           fp16_res=fp16_res)
+
+        return {**mapping_results, **synthesis_results}
+
+
+class MappingNetwork(nn.Module):
+    """Implements the latent space mapping network.
+
+    Basically, this network executes several dense layers in sequence, and the
+    label embedding if needed.
+    """
+
+    def __init__(self,
+                 input_dim,
+                 output_dim,
+                 num_outputs,
+                 repeat_output,
+                 normalize_input,
+                 num_layers,
+                 hidden_dim,
+                 use_wscale,
+                 wscale_gain,
+                 lr_mul,
+                 label_dim,
+                 embedding_dim,
+                 embedding_bias,
+                 embedding_use_wscale,
+                 embedding_wscale_gian,
+                 embedding_lr_mul,
+                 normalize_embedding,
+                 normalize_embedding_latent,
+                 eps):
+        super().__init__()
+
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+        self.num_outputs = num_outputs
+        self.repeat_output = repeat_output
+        self.normalize_input = normalize_input
+        self.num_layers = num_layers
+        self.hidden_dim = hidden_dim
+        self.use_wscale = use_wscale
+        self.wscale_gain = wscale_gain
+        self.lr_mul = lr_mul
+        self.label_dim = label_dim
+        self.embedding_dim = embedding_dim
+        self.embedding_bias = embedding_bias
+        self.embedding_use_wscale = embedding_use_wscale
+        self.embedding_wscale_gian = embedding_wscale_gian
+        self.embedding_lr_mul = embedding_lr_mul
+        self.normalize_embedding = normalize_embedding
+        self.normalize_embedding_latent = normalize_embedding_latent
+        self.eps = eps
+
+        self.pth_to_tf_var_mapping = {}
+
+        self.norm = PixelNormLayer(dim=1, eps=eps)
+
+        if self.label_dim > 0:
+            input_dim = input_dim + embedding_dim
+            self.embedding = DenseLayer(in_channels=label_dim,
+                                        out_channels=embedding_dim,
+                                        add_bias=embedding_bias,
+                                        init_bias=0.0,
+                                        use_wscale=embedding_use_wscale,
+                                        wscale_gain=embedding_wscale_gian,
+                                        lr_mul=embedding_lr_mul,
+                                        activation_type='linear')
+            self.pth_to_tf_var_mapping['embedding.weight'] = 'LabelEmbed/weight'
+            if self.embedding_bias:
+                self.pth_to_tf_var_mapping['embedding.bias'] = 'LabelEmbed/bias'
+
+        if num_outputs is not None and not repeat_output:
+            output_dim = output_dim * num_outputs
+        for i in range(num_layers):
+            in_channels = (input_dim if i == 0 else hidden_dim)
+            out_channels = (output_dim if i == (num_layers - 1) else hidden_dim)
+            self.add_module(f'dense{i}',
+                            DenseLayer(in_channels=in_channels,
+                                       out_channels=out_channels,
+                                       add_bias=True,
+                                       init_bias=0.0,
+                                       use_wscale=use_wscale,
+                                       wscale_gain=wscale_gain,
+                                       lr_mul=lr_mul,
+                                       activation_type='lrelu'))
+            self.pth_to_tf_var_mapping[f'dense{i}.weight'] = f'Dense{i}/weight'
+            self.pth_to_tf_var_mapping[f'dense{i}.bias'] = f'Dense{i}/bias'
+
+    def forward(self, z, label=None, impl='cuda'):
+        if z.ndim != 2 or z.shape[1] != self.input_dim:
+            raise ValueError(f'Input latent code should be with shape '
+                             f'[batch_size, input_dim], where '
+                             f'`input_dim` equals to {self.input_dim}!\n'
+                             f'But `{z.shape}` is received!')
+        if self.normalize_input:
+            z = self.norm(z)
+
+        if self.label_dim > 0:
+            if label is None:
+                raise ValueError(f'Model requires an additional label '
+                                 f'(with dimension {self.label_dim}) as input, '
+                                 f'but no label is received!')
+            if label.ndim != 2 or label.shape != (z.shape[0], self.label_dim):
+                raise ValueError(f'Input label should be with shape '
+                                 f'[batch_size, label_dim], where '
+                                 f'`batch_size` equals to that of '
+                                 f'latent codes ({z.shape[0]}) and '
+                                 f'`label_dim` equals to {self.label_dim}!\n'
+                                 f'But `{label.shape}` is received!')
+            label = label.to(dtype=torch.float32)
+            embedding = self.embedding(label, impl=impl)
+            if self.normalize_embedding:
+                embedding = self.norm(embedding)
+            w = torch.cat((z, embedding), dim=1)
+        else:
+            w = z
+
+        if self.label_dim > 0 and self.normalize_embedding_latent:
+            w = self.norm(w)
+
+        for i in range(self.num_layers):
+            w = getattr(self, f'dense{i}')(w, impl=impl)
+
+        wp = None
+        if self.num_outputs is not None:
+            if self.repeat_output:
+                wp = w.unsqueeze(1).repeat((1, self.num_outputs, 1))
+            else:
+                wp = w.reshape(-1, self.num_outputs, self.output_dim)
+
+        results = {
+            'z': z,
+            'label': label,
+            'w': w,
+            'wp': wp,
+        }
+        if self.label_dim > 0:
+            results['embedding'] = embedding
+        return results
+
+
+class SynthesisNetwork(nn.Module):
+    """Implements the image synthesis network.
+
+    Basically, this network executes several convolutional layers in sequence.
+    """
+
+    def __init__(self,
+                 resolution,
+                 init_res,
+                 w_dim,
+                 image_channels,
+                 final_tanh,
+                 const_input,
+                 architecture,
+                 demodulate,
+                 use_wscale,
+                 wscale_gain,
+                 lr_mul,
+                 noise_type,
+                 fmaps_base,
+                 fmaps_max,
+                 filter_kernel,
+                 conv_clamp,
+                 eps):
+        super().__init__()
+
+        self.init_res = init_res
+        self.init_res_log2 = int(np.log2(init_res))
+        self.resolution = resolution
+        self.final_res_log2 = int(np.log2(resolution))
+        self.w_dim = w_dim
+        self.image_channels = image_channels
+        self.final_tanh = final_tanh
+        self.const_input = const_input
+        self.architecture = architecture.lower()
+        self.demodulate = demodulate
+        self.use_wscale = use_wscale
+        self.wscale_gain = wscale_gain
+        self.lr_mul = lr_mul
+        self.noise_type = noise_type.lower()
+        self.fmaps_base = fmaps_base
+        self.fmaps_max = fmaps_max
+        self.filter_kernel = filter_kernel
+        self.conv_clamp = conv_clamp
+        self.eps = eps
+
+        self.num_layers = (self.final_res_log2 - self.init_res_log2 + 1) * 2
+
+        self.pth_to_tf_var_mapping = {}
+
+        for res_log2 in range(self.init_res_log2, self.final_res_log2 + 1):
+            res = 2 ** res_log2
+            in_channels = self.get_nf(res // 2)
+            out_channels = self.get_nf(res)
+            block_idx = res_log2 - self.init_res_log2
+
+            # Early layer.
+            if res == init_res:
+                if self.const_input:
+                    self.add_module('early_layer',
+                                    InputLayer(init_res=res,
+                                               channels=out_channels))
+                    self.pth_to_tf_var_mapping['early_layer.const'] = (
+                        f'{res}x{res}/Const/const')
+                else:
+                    channels = out_channels * res * res
+                    self.add_module('early_layer',
+                                    DenseLayer(in_channels=w_dim,
+                                               out_channels=channels,
+                                               add_bias=True,
+                                               init_bias=0.0,
+                                               use_wscale=use_wscale,
+                                               wscale_gain=wscale_gain,
+                                               lr_mul=lr_mul,
+                                               activation_type='lrelu'))
+                    self.pth_to_tf_var_mapping['early_layer.weight'] = (
+                        f'{res}x{res}/Dense/weight')
+                    self.pth_to_tf_var_mapping['early_layer.bias'] = (
+                        f'{res}x{res}/Dense/bias')
+            else:
+                # Residual branch (kernel 1x1) with upsampling, without bias,
+                # with linear activation.
+                if self.architecture == 'resnet':
+                    layer_name = f'residual{block_idx}'
+                    self.add_module(layer_name,
+                                    ConvLayer(in_channels=in_channels,
+                                              out_channels=out_channels,
+                                              kernel_size=1,
+                                              add_bias=False,
+                                              scale_factor=2,
+                                              filter_kernel=filter_kernel,
+                                              use_wscale=use_wscale,
+                                              wscale_gain=wscale_gain,
+                                              lr_mul=lr_mul,
+                                              activation_type='linear',
+                                              conv_clamp=None))
+                    self.pth_to_tf_var_mapping[f'{layer_name}.weight'] = (
+                        f'{res}x{res}/Skip/weight')
+
+                # First layer (kernel 3x3) with upsampling.
+                layer_name = f'layer{2 * block_idx - 1}'
+                self.add_module(layer_name,
+                                ModulateConvLayer(in_channels=in_channels,
+                                                  out_channels=out_channels,
+                                                  resolution=res,
+                                                  w_dim=w_dim,
+                                                  kernel_size=3,
+                                                  add_bias=True,
+                                                  scale_factor=2,
+                                                  filter_kernel=filter_kernel,
+                                                  demodulate=demodulate,
+                                                  use_wscale=use_wscale,
+                                                  wscale_gain=wscale_gain,
+                                                  lr_mul=lr_mul,
+                                                  noise_type=noise_type,
+                                                  activation_type='lrelu',
+                                                  conv_clamp=conv_clamp,
+                                                  eps=eps))
+                self.pth_to_tf_var_mapping[f'{layer_name}.weight'] = (
+                    f'{res}x{res}/Conv0_up/weight')
+                self.pth_to_tf_var_mapping[f'{layer_name}.bias'] = (
+                    f'{res}x{res}/Conv0_up/bias')
+                self.pth_to_tf_var_mapping[f'{layer_name}.style.weight'] = (
+                    f'{res}x{res}/Conv0_up/mod_weight')
+                self.pth_to_tf_var_mapping[f'{layer_name}.style.bias'] = (
+                    f'{res}x{res}/Conv0_up/mod_bias')
+                self.pth_to_tf_var_mapping[f'{layer_name}.noise_strength'] = (
+                    f'{res}x{res}/Conv0_up/noise_strength')
+                self.pth_to_tf_var_mapping[f'{layer_name}.noise'] = (
+                    f'noise{2 * block_idx - 1}')
+
+            # Second layer (kernel 3x3) without upsampling.
+            layer_name = f'layer{2 * block_idx}'
+            self.add_module(layer_name,
+                            ModulateConvLayer(in_channels=out_channels,
+                                              out_channels=out_channels,
+                                              resolution=res,
+                                              w_dim=w_dim,
+                                              kernel_size=3,
+                                              add_bias=True,
+                                              scale_factor=1,
+                                              filter_kernel=None,
+                                              demodulate=demodulate,
+                                              use_wscale=use_wscale,
+                                              wscale_gain=wscale_gain,
+                                              lr_mul=lr_mul,
+                                              noise_type=noise_type,
+                                              activation_type='lrelu',
+                                              conv_clamp=conv_clamp,
+                                              eps=eps))
+            tf_layer_name = 'Conv' if res == self.init_res else 'Conv1'
+            self.pth_to_tf_var_mapping[f'{layer_name}.weight'] = (
+                f'{res}x{res}/{tf_layer_name}/weight')
+            self.pth_to_tf_var_mapping[f'{layer_name}.bias'] = (
+                f'{res}x{res}/{tf_layer_name}/bias')
+            self.pth_to_tf_var_mapping[f'{layer_name}.style.weight'] = (
+                f'{res}x{res}/{tf_layer_name}/mod_weight')
+            self.pth_to_tf_var_mapping[f'{layer_name}.style.bias'] = (
+                f'{res}x{res}/{tf_layer_name}/mod_bias')
+            self.pth_to_tf_var_mapping[f'{layer_name}.noise_strength'] = (
+                f'{res}x{res}/{tf_layer_name}/noise_strength')
+            self.pth_to_tf_var_mapping[f'{layer_name}.noise'] = (
+                f'noise{2 * block_idx}')
+
+            # Output convolution layer for each resolution (if needed).
+            if res_log2 == self.final_res_log2 or self.architecture == 'skip':
+                layer_name = f'output{block_idx}'
+                self.add_module(layer_name,
+                                ModulateConvLayer(in_channels=out_channels,
+                                                  out_channels=image_channels,
+                                                  resolution=res,
+                                                  w_dim=w_dim,
+                                                  kernel_size=1,
+                                                  add_bias=True,
+                                                  scale_factor=1,
+                                                  filter_kernel=None,
+                                                  demodulate=False,
+                                                  use_wscale=use_wscale,
+                                                  wscale_gain=wscale_gain,
+                                                  lr_mul=lr_mul,
+                                                  noise_type='none',
+                                                  activation_type='linear',
+                                                  conv_clamp=conv_clamp,
+                                                  eps=eps))
+                self.pth_to_tf_var_mapping[f'{layer_name}.weight'] = (
+                    f'{res}x{res}/ToRGB/weight')
+                self.pth_to_tf_var_mapping[f'{layer_name}.bias'] = (
+                    f'{res}x{res}/ToRGB/bias')
+                self.pth_to_tf_var_mapping[f'{layer_name}.style.weight'] = (
+                    f'{res}x{res}/ToRGB/mod_weight')
+                self.pth_to_tf_var_mapping[f'{layer_name}.style.bias'] = (
+                    f'{res}x{res}/ToRGB/mod_bias')
+
+        # Used for upsampling output images for each resolution block for sum.
+        if self.architecture == 'skip':
+            self.register_buffer(
+                'filter', upfirdn2d.setup_filter(filter_kernel))
+
+    def get_nf(self, res):
+        """Gets number of feature maps according to the given resolution."""
+        return min(self.fmaps_base // res, self.fmaps_max)
+
+    def set_space_of_latent(self, space_of_latent):
+        """Sets the space to which the latent code belong.
+
+        This function is particularly used for choosing how to inject the latent
+        code into the convolutional layers. The original generator will take a
+        W-Space code and apply it for style modulation after an affine
+        transformation. But, sometimes, it may need to directly feed an already
+        affine-transformed code into the convolutional layer, e.g., when
+        training an encoder for GAN inversion. We term the transformed space as
+        Style Space (or Y-Space). This function is designed to tell the
+        convolutional layers how to use the input code.
+
+        Args:
+            space_of_latent: The space to which the latent code belong. Case
+                insensitive. Support `W` and `Y`.
+        """
+        space_of_latent = space_of_latent.upper()
+        for module in self.modules():
+            if isinstance(module, ModulateConvLayer):
+                setattr(module, 'space_of_latent', space_of_latent)
+
+    def forward(self,
+                wp,
+                noise_mode='const',
+                fused_modulate=False,
+                fp16_res=None,
+                impl='cuda'):
+        results = {'wp': wp}
+
+        if self.const_input:
+            x = self.early_layer(wp[:, 0])
+        else:
+            x = self.early_layer(wp[:, 0], impl=impl)
+
+        # Cast to `torch.float16` if needed.
+        if fp16_res is not None and self.init_res >= fp16_res:
+            x = x.to(torch.float16)
+
+        if self.architecture == 'origin':
+            for layer_idx in range(self.num_layers - 1):
+                layer = getattr(self, f'layer{layer_idx}')
+                x, style = layer(x,
+                                 wp[:, layer_idx],
+                                 noise_mode=noise_mode,
+                                 fused_modulate=fused_modulate,
+                                 impl=impl)
+                results[f'style{layer_idx}'] = style
+
+                # Cast to `torch.float16` if needed.
+                if layer_idx % 2 == 0 and layer_idx != self.num_layers - 2:
+                    res = self.init_res * (2 ** (layer_idx // 2))
+                    if fp16_res is not None and res * 2 >= fp16_res:
+                        x = x.to(torch.float16)
+                    else:
+                        x = x.to(torch.float32)
+            output_layer = getattr(self, f'output{layer_idx // 2}')
+            image, style = output_layer(x,
+                                        wp[:, layer_idx + 1],
+                                        fused_modulate=fused_modulate,
+                                        impl=impl)
+            image = image.to(torch.float32)
+            results[f'output_style{layer_idx // 2}'] = style
+
+        elif self.architecture == 'skip':
+            for layer_idx in range(self.num_layers - 1):
+                layer = getattr(self, f'layer{layer_idx}')
+                x, style = layer(x,
+                                 wp[:, layer_idx],
+                                 noise_mode=noise_mode,
+                                 fused_modulate=fused_modulate,
+                                 impl=impl)
+                results[f'style{layer_idx}'] = style
+                if layer_idx % 2 == 0:
+                    output_layer = getattr(self, f'output{layer_idx // 2}')
+                    y, style = output_layer(x,
+                                            wp[:, layer_idx + 1],
+                                            fused_modulate=fused_modulate,
+                                            impl=impl)
+                    results[f'output_style{layer_idx // 2}'] = style
+                    if layer_idx == 0:
+                        image = y.to(torch.float32)
+                    else:
+                        image = y.to(torch.float32) + upfirdn2d.upsample2d(
+                            image, self.filter, impl=impl)
+
+                    # Cast to `torch.float16` if needed.
+                    if layer_idx != self.num_layers - 2:
+                        res = self.init_res * (2 ** (layer_idx // 2))
+                        if fp16_res is not None and res * 2 >= fp16_res:
+                            x = x.to(torch.float16)
+                        else:
+                            x = x.to(torch.float32)
+
+        elif self.architecture == 'resnet':
+            x, style = self.layer0(x,
+                                   wp[:, 0],
+                                   noise_mode=noise_mode,
+                                   fused_modulate=fused_modulate,
+                                   impl=impl)
+            results['style0'] = style
+            for layer_idx in range(1, self.num_layers - 1, 2):
+                # Cast to `torch.float16` if needed.
+                if layer_idx % 2 == 1:
+                    res = self.init_res * (2 ** (layer_idx // 2))
+                    if fp16_res is not None and res * 2 >= fp16_res:
+                        x = x.to(torch.float16)
+                    else:
+                        x = x.to(torch.float32)
+
+                skip_layer = getattr(self, f'residual{layer_idx // 2 + 1}')
+                residual = skip_layer(x, runtime_gain=np.sqrt(0.5), impl=impl)
+                layer = getattr(self, f'layer{layer_idx}')
+                x, style = layer(x,
+                                 wp[:, layer_idx],
+                                 noise_mode=noise_mode,
+                                 fused_modulate=fused_modulate,
+                                 impl=impl)
+                results[f'style{layer_idx}'] = style
+                layer = getattr(self, f'layer{layer_idx + 1}')
+                x, style = layer(x,
+                                 wp[:, layer_idx + 1],
+                                 runtime_gain=np.sqrt(0.5),
+                                 noise_mode=noise_mode,
+                                 fused_modulate=fused_modulate,
+                                 impl=impl)
+                results[f'style{layer_idx + 1}'] = style
+                x = x + residual
+            output_layer = getattr(self, f'output{layer_idx // 2 + 1}')
+            image, style = output_layer(x,
+                                        wp[:, layer_idx + 2],
+                                        fused_modulate=fused_modulate,
+                                        impl=impl)
+            image = image.to(torch.float32)
+            results[f'output_style{layer_idx // 2}'] = style
+
+        if self.final_tanh:
+            image = torch.tanh(image)
+        results['image'] = image
+        return results
+
+
+class PixelNormLayer(nn.Module):
+    """Implements pixel-wise feature vector normalization layer."""
+
+    def __init__(self, dim, eps):
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+
+    def extra_repr(self):
+        return f'dim={self.dim}, epsilon={self.eps}'
+
+    def forward(self, x):
+        scale = (x.square().mean(dim=self.dim, keepdim=True) + self.eps).rsqrt()
+        return x * scale
+
+
+class InputLayer(nn.Module):
+    """Implements the input layer to start convolution with.
+
+    Basically, this block starts from a const input, which is with shape
+    `(channels, init_res, init_res)`.
+    """
+
+    def __init__(self, init_res, channels):
+        super().__init__()
+        self.const = nn.Parameter(torch.randn(1, channels, init_res, init_res))
+
+    def forward(self, w):
+        x = self.const.repeat(w.shape[0], 1, 1, 1)
+        return x
+
+
+class ConvLayer(nn.Module):
+    """Implements the convolutional layer.
+
+    If upsampling is needed (i.e., `scale_factor = 2`), the feature map will
+    be filtered with `filter_kernel` after convolution. This layer will only be
+    used for skip connection in `resnet` architecture.
+    """
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 add_bias,
+                 scale_factor,
+                 filter_kernel,
+                 use_wscale,
+                 wscale_gain,
+                 lr_mul,
+                 activation_type,
+                 conv_clamp):
+        """Initializes with layer settings.
+
+        Args:
+            in_channels: Number of channels of the input tensor.
+            out_channels: Number of channels of the output tensor.
+            kernel_size: Size of the convolutional kernels.
+            add_bias: Whether to add bias onto the convolutional result.
+            scale_factor: Scale factor for upsampling.
+            filter_kernel: Kernel used for filtering.
+            use_wscale: Whether to use weight scaling.
+            wscale_gain: Gain factor for weight scaling.
+            lr_mul: Learning multiplier for both weight and bias.
+            activation_type: Type of activation.
+            conv_clamp: A threshold to clamp the output of convolution layers to
+                avoid overflow under FP16 training.
+        """
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        self.add_bias = add_bias
+        self.scale_factor = scale_factor
+        self.filter_kernel = filter_kernel
+        self.use_wscale = use_wscale
+        self.wscale_gain = wscale_gain
+        self.lr_mul = lr_mul
+        self.activation_type = activation_type
+        self.conv_clamp = conv_clamp
+
+        weight_shape = (out_channels, in_channels, kernel_size, kernel_size)
+        fan_in = kernel_size * kernel_size * in_channels
+        wscale = wscale_gain / np.sqrt(fan_in)
+        if use_wscale:
+            self.weight = nn.Parameter(torch.randn(*weight_shape) / lr_mul)
+            self.wscale = wscale * lr_mul
+        else:
+            self.weight = nn.Parameter(
+                torch.randn(*weight_shape) * wscale / lr_mul)
+            self.wscale = lr_mul
+
+        if add_bias:
+            self.bias = nn.Parameter(torch.zeros(out_channels))
+            self.bscale = lr_mul
+        else:
+            self.bias = None
+        self.act_gain = bias_act.activation_funcs[activation_type].def_gain
+
+        if scale_factor > 1:
+            assert filter_kernel is not None
+            self.register_buffer(
+                'filter', upfirdn2d.setup_filter(filter_kernel))
+            fh, fw = self.filter.shape
+            self.filter_padding = (
+                kernel_size // 2 + (fw + scale_factor - 1) // 2,
+                kernel_size // 2 + (fw - scale_factor) // 2,
+                kernel_size // 2 + (fh + scale_factor - 1) // 2,
+                kernel_size // 2 + (fh - scale_factor) // 2)
+
+    def extra_repr(self):
+        return (f'in_ch={self.in_channels}, '
+                f'out_ch={self.out_channels}, '
+                f'ksize={self.kernel_size}, '
+                f'wscale_gain={self.wscale_gain:.3f}, '
+                f'bias={self.add_bias}, '
+                f'lr_mul={self.lr_mul:.3f}, '
+                f'upsample={self.scale_factor}, '
+                f'upsample_filter={self.filter_kernel}, '
+                f'act={self.activation_type}, '
+                f'clamp={self.conv_clamp}')
+
+    def forward(self, x, runtime_gain=1.0, impl='cuda'):
+        dtype = x.dtype
+
+        weight = self.weight
+        if self.wscale != 1.0:
+            weight = weight * self.wscale
+        bias = None
+        if self.bias is not None:
+            bias = self.bias.to(dtype)
+            if self.bscale != 1.0:
+                bias = bias * self.bscale
+
+        if self.scale_factor == 1:  # Native convolution without upsampling.
+            padding = self.kernel_size // 2
+            x = conv2d_gradfix.conv2d(
+                x, weight.to(dtype), stride=1, padding=padding, impl=impl)
+        else:  # Convolution with upsampling.
+            up = self.scale_factor
+            f = self.filter
+            # When kernel size = 1, use filtering function for upsampling.
+            if self.kernel_size == 1:
+                padding = self.filter_padding
+                x = conv2d_gradfix.conv2d(
+                    x, weight.to(dtype), stride=1, padding=0, impl=impl)
+                x = upfirdn2d.upfirdn2d(
+                    x, f, up=up, padding=padding, gain=up ** 2, impl=impl)
+            # When kernel size != 1, use transpose convolution for upsampling.
+            else:
+                # Following codes are borrowed from
+                # https://github.com/NVlabs/stylegan2-ada-pytorch
+                px0, px1, py0, py1 = self.filter_padding
+                kh, kw = weight.shape[2:]
+                px0 = px0 - (kw - 1)
+                px1 = px1 - (kw - up)
+                py0 = py0 - (kh - 1)
+                py1 = py1 - (kh - up)
+                pxt = max(min(-px0, -px1), 0)
+                pyt = max(min(-py0, -py1), 0)
+                weight = weight.transpose(0, 1)
+                padding = (pyt, pxt)
+                x = conv2d_gradfix.conv_transpose2d(
+                    x, weight.to(dtype), stride=up, padding=padding, impl=impl)
+                padding = (px0 + pxt, px1 + pxt, py0 + pyt, py1 + pyt)
+                x = upfirdn2d.upfirdn2d(
+                    x, f, up=1, padding=padding, gain=up ** 2, impl=impl)
+
+        act_gain = self.act_gain * runtime_gain
+        act_clamp = None
+        if self.conv_clamp is not None:
+            act_clamp = self.conv_clamp * runtime_gain
+        x = bias_act.bias_act(x, bias,
+                              act=self.activation_type,
+                              gain=act_gain,
+                              clamp=act_clamp,
+                              impl=impl)
+
+        assert x.dtype == dtype
+        return x
+
+
+class ModulateConvLayer(nn.Module):
+    """Implements the convolutional layer with style modulation."""
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 resolution,
+                 w_dim,
+                 kernel_size,
+                 add_bias,
+                 scale_factor,
+                 filter_kernel,
+                 demodulate,
+                 use_wscale,
+                 wscale_gain,
+                 lr_mul,
+                 noise_type,
+                 activation_type,
+                 conv_clamp,
+                 eps):
+        """Initializes with layer settings.
+
+        Args:
+            in_channels: Number of channels of the input tensor.
+            out_channels: Number of channels of the output tensor.
+            resolution: Resolution of the output tensor.
+            w_dim: Dimension of W space for style modulation.
+            kernel_size: Size of the convolutional kernels.
+            add_bias: Whether to add bias onto the convolutional result.
+            scale_factor: Scale factor for upsampling.
+            filter_kernel: Kernel used for filtering.
+            demodulate: Whether to perform style demodulation.
+            use_wscale: Whether to use weight scaling.
+            wscale_gain: Gain factor for weight scaling.
+            lr_mul: Learning multiplier for both weight and bias.
+            noise_type: Type of noise added to the feature map after the
+                convolution (if needed). Support `none`, `spatial` and
+                `channel`.
+            activation_type: Type of activation.
+            conv_clamp: A threshold to clamp the output of convolution layers to
+                avoid overflow under FP16 training.
+            eps: A small value to avoid divide overflow.
+        """
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.resolution = resolution
+        self.w_dim = w_dim
+        self.kernel_size = kernel_size
+        self.add_bias = add_bias
+        self.scale_factor = scale_factor
+        self.filter_kernel = filter_kernel
+        self.demodulate = demodulate
+        self.use_wscale = use_wscale
+        self.wscale_gain = wscale_gain
+        self.lr_mul = lr_mul
+        self.noise_type = noise_type.lower()
+        self.activation_type = activation_type
+        self.conv_clamp = conv_clamp
+        self.eps = eps
+
+        self.space_of_latent = 'W'
+
+        # Set up weight.
+        weight_shape = (out_channels, in_channels, kernel_size, kernel_size)
+        fan_in = kernel_size * kernel_size * in_channels
+        wscale = wscale_gain / np.sqrt(fan_in)
+        if use_wscale:
+            self.weight = nn.Parameter(torch.randn(*weight_shape) / lr_mul)
+            self.wscale = wscale * lr_mul
+        else:
+            self.weight = nn.Parameter(
+                torch.randn(*weight_shape) * wscale / lr_mul)
+            self.wscale = lr_mul
+
+        # Set up bias.
+        if add_bias:
+            self.bias = nn.Parameter(torch.zeros(out_channels))
+            self.bscale = lr_mul
+        else:
+            self.bias = None
+        self.act_gain = bias_act.activation_funcs[activation_type].def_gain
+
+        # Set up style.
+        self.style = DenseLayer(in_channels=w_dim,
+                                out_channels=in_channels,
+                                add_bias=True,
+                                init_bias=1.0,
+                                use_wscale=use_wscale,
+                                wscale_gain=wscale_gain,
+                                lr_mul=lr_mul,
+                                activation_type='linear')
+
+        # Set up noise.
+        if self.noise_type != 'none':
+            self.noise_strength = nn.Parameter(torch.zeros(()))
+            if self.noise_type == 'spatial':
+                self.register_buffer(
+                    'noise', torch.randn(1, 1, resolution, resolution))
+            elif self.noise_type == 'channel':
+                self.register_buffer(
+                    'noise', torch.randn(1, out_channels, 1, 1))
+            else:
+                raise NotImplementedError(f'Not implemented noise type: '
+                                          f'`{self.noise_type}`!')
+
+        if scale_factor > 1:
+            assert filter_kernel is not None
+            self.register_buffer(
+                'filter', upfirdn2d.setup_filter(filter_kernel))
+            fh, fw = self.filter.shape
+            self.filter_padding = (
+                kernel_size // 2 + (fw + scale_factor - 1) // 2,
+                kernel_size // 2 + (fw - scale_factor) // 2,
+                kernel_size // 2 + (fh + scale_factor - 1) // 2,
+                kernel_size // 2 + (fh - scale_factor) // 2)
+
+    def extra_repr(self):
+        return (f'in_ch={self.in_channels}, '
+                f'out_ch={self.out_channels}, '
+                f'ksize={self.kernel_size}, '
+                f'wscale_gain={self.wscale_gain:.3f}, '
+                f'bias={self.add_bias}, '
+                f'lr_mul={self.lr_mul:.3f}, '
+                f'upsample={self.scale_factor}, '
+                f'upsample_filter={self.filter_kernel}, '
+                f'demodulate={self.demodulate}, '
+                f'noise_type={self.noise_type}, '
+                f'act={self.activation_type}, '
+                f'clamp={self.conv_clamp}')
+
+    def forward_style(self, w, impl='cuda'):
+        """Gets style code from the given input.
+
+        More specifically, if the input is from W-Space, it will be projected by
+        an affine transformation. If it is from the Style Space (Y-Space), no
+        operation is required.
+
+        NOTE: For codes from Y-Space, we use slicing to make sure the dimension
+        is correct, in case that the code is padded before fed into this layer.
+        """
+        space_of_latent = self.space_of_latent.upper()
+        if space_of_latent == 'W':
+            if w.ndim != 2 or w.shape[1] != self.w_dim:
+                raise ValueError(f'The input tensor should be with shape '
+                                 f'[batch_size, w_dim], where '
+                                 f'`w_dim` equals to {self.w_dim}!\n'
+                                 f'But `{w.shape}` is received!')
+            style = self.style(w, impl=impl)
+        elif space_of_latent == 'Y':
+            if w.ndim != 2 or w.shape[1] < self.in_channels:
+                raise ValueError(f'The input tensor should be with shape '
+                                 f'[batch_size, y_dim], where '
+                                 f'`y_dim` equals to {self.in_channels}!\n'
+                                 f'But `{w.shape}` is received!')
+            style = w[:, :self.in_channels]
+        else:
+            raise NotImplementedError(f'Not implemented `space_of_latent`: '
+                                      f'`{space_of_latent}`!')
+        return style
+
+    def forward(self,
+                x,
+                w,
+                runtime_gain=1.0,
+                noise_mode='const',
+                fused_modulate=False,
+                impl='cuda'):
+        dtype = x.dtype
+        N, C, H, W = x.shape
+
+        fused_modulate = (fused_modulate and
+                          not self.training and
+                          (dtype == torch.float32 or N == 1))
+
+        weight = self.weight
+        out_ch, in_ch, kh, kw = weight.shape
+        assert in_ch == C
+
+        # Affine on `w`.
+        style = self.forward_style(w, impl=impl)
+        if not self.demodulate:
+            _style = style * self.wscale  # Equivalent to scaling weight.
+        else:
+            _style = style
+
+        # Prepare noise.
+        noise = None
+        noise_mode = noise_mode.lower()
+        if self.noise_type != 'none' and noise_mode != 'none':
+            if noise_mode == 'random':
+                noise = torch.randn((N, *self.noise.shape[1:]), device=x.device)
+            elif noise_mode == 'const':
+                noise = self.noise
+            else:
+                raise ValueError(f'Unknown noise mode `{noise_mode}`!')
+            noise = (noise * self.noise_strength).to(dtype)
+
+        # Pre-normalize inputs to avoid FP16 overflow.
+        if dtype == torch.float16 and self.demodulate:
+            weight_max = weight.norm(float('inf'), dim=(1, 2, 3), keepdim=True)
+            weight = weight * (self.wscale / weight_max)
+            style_max = _style.norm(float('inf'), dim=1, keepdim=True)
+            _style = _style / style_max
+
+        if self.demodulate or fused_modulate:
+            _weight = weight.unsqueeze(0)
+            _weight = _weight * _style.reshape(N, 1, in_ch, 1, 1)
+        if self.demodulate:
+            decoef = (_weight.square().sum(dim=(2, 3, 4)) + self.eps).rsqrt()
+        if self.demodulate and fused_modulate:
+            _weight = _weight * decoef.reshape(N, out_ch, 1, 1, 1)
+
+        if not fused_modulate:
+            x = x * _style.to(dtype).reshape(N, in_ch, 1, 1)
+            w = weight.to(dtype)
+            groups = 1
+        else:  # Use group convolution to fuse style modulation and convolution.
+            x = x.reshape(1, N * in_ch, H, W)
+            w = _weight.reshape(N * out_ch, in_ch, kh, kw).to(dtype)
+            groups = N
+
+        if self.scale_factor == 1:  # Native convolution without upsampling.
+            up = 1
+            padding = self.kernel_size // 2
+            x = conv2d_gradfix.conv2d(
+                x, w, stride=1, padding=padding, groups=groups, impl=impl)
+        else:  # Convolution with upsampling.
+            up = self.scale_factor
+            f = self.filter
+            # When kernel size = 1, use filtering function for upsampling.
+            if self.kernel_size == 1:
+                padding = self.filter_padding
+                x = conv2d_gradfix.conv2d(
+                    x, w, stride=1, padding=0, groups=groups, impl=impl)
+                x = upfirdn2d.upfirdn2d(
+                    x, f, up=up, padding=padding, gain=up ** 2, impl=impl)
+            # When kernel size != 1, use stride convolution for upsampling.
+            else:
+                # Following codes are borrowed from
+                # https://github.com/NVlabs/stylegan2-ada-pytorch
+                px0, px1, py0, py1 = self.filter_padding
+                px0 = px0 - (kw - 1)
+                px1 = px1 - (kw - up)
+                py0 = py0 - (kh - 1)
+                py1 = py1 - (kh - up)
+                pxt = max(min(-px0, -px1), 0)
+                pyt = max(min(-py0, -py1), 0)
+                if groups == 1:
+                    w = w.transpose(0, 1)
+                else:
+                    w = w.reshape(N, out_ch, in_ch, kh, kw)
+                    w = w.transpose(1, 2)
+                    w = w.reshape(N * in_ch, out_ch, kh, kw)
+                padding = (pyt, pxt)
+                x = conv2d_gradfix.conv_transpose2d(
+                    x, w, stride=up, padding=padding, groups=groups, impl=impl)
+                padding = (px0 + pxt, px1 + pxt, py0 + pyt, py1 + pyt)
+                x = upfirdn2d.upfirdn2d(
+                    x, f, up=1, padding=padding, gain=up ** 2, impl=impl)
+
+        if not fused_modulate:
+            if self.demodulate:
+                decoef = decoef.to(dtype).reshape(N, out_ch, 1, 1)
+            if self.demodulate and noise is not None:
+                x = fma.fma(x, decoef, noise, impl=impl)
+            else:
+                if self.demodulate:
+                    x = x * decoef
+                if noise is not None:
+                    x = x + noise
+        else:
+            x = x.reshape(N, out_ch, H * up, W * up)
+            if noise is not None:
+                x = x + noise
+
+        bias = None
+        if self.bias is not None:
+            bias = self.bias.to(dtype)
+            if self.bscale != 1.0:
+                bias = bias * self.bscale
+
+        if self.activation_type == 'linear':  # Shortcut for output layer.
+            x = bias_act.bias_act(
+                x, bias, act='linear', clamp=self.conv_clamp, impl=impl)
+        else:
+            act_gain = self.act_gain * runtime_gain
+            act_clamp = None
+            if self.conv_clamp is not None:
+                act_clamp = self.conv_clamp * runtime_gain
+            x = bias_act.bias_act(x, bias,
+                                  act=self.activation_type,
+                                  gain=act_gain,
+                                  clamp=act_clamp,
+                                  impl=impl)
+
+        assert x.dtype == dtype
+        assert style.dtype == torch.float32
+        return x, style
+
+
+class DenseLayer(nn.Module):
+    """Implements the dense layer."""
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 add_bias,
+                 init_bias,
+                 use_wscale,
+                 wscale_gain,
+                 lr_mul,
+                 activation_type):
+        """Initializes with layer settings.
+
+        Args:
+            in_channels: Number of channels of the input tensor.
+            out_channels: Number of channels of the output tensor.
+            add_bias: Whether to add bias onto the fully-connected result.
+            init_bias: The initial bias value before training.
+            use_wscale: Whether to use weight scaling.
+            wscale_gain: Gain factor for weight scaling.
+            lr_mul: Learning multiplier for both weight and bias.
+            activation_type: Type of activation.
+        """
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.add_bias = add_bias
+        self.init_bias = init_bias
+        self.use_wscale = use_wscale
+        self.wscale_gain = wscale_gain
+        self.lr_mul = lr_mul
+        self.activation_type = activation_type
+
+        weight_shape = (out_channels, in_channels)
+        wscale = wscale_gain / np.sqrt(in_channels)
+        if use_wscale:
+            self.weight = nn.Parameter(torch.randn(*weight_shape) / lr_mul)
+            self.wscale = wscale * lr_mul
+        else:
+            self.weight = nn.Parameter(
+                torch.randn(*weight_shape) * wscale / lr_mul)
+            self.wscale = lr_mul
+
+        if add_bias:
+            init_bias = np.float32(init_bias) / lr_mul
+            self.bias = nn.Parameter(torch.full([out_channels], init_bias))
+            self.bscale = lr_mul
+        else:
+            self.bias = None
+
+    def extra_repr(self):
+        return (f'in_ch={self.in_channels}, '
+                f'out_ch={self.out_channels}, '
+                f'wscale_gain={self.wscale_gain:.3f}, '
+                f'bias={self.add_bias}, '
+                f'init_bias={self.init_bias}, '
+                f'lr_mul={self.lr_mul:.3f}, '
+                f'act={self.activation_type}')
+
+    def forward(self, x, impl='cuda'):
+        dtype = x.dtype
+
+        if x.ndim != 2:
+            x = x.flatten(start_dim=1)
+
+        weight = self.weight.to(dtype) * self.wscale
+        bias = None
+        if self.bias is not None:
+            bias = self.bias.to(dtype)
+            if self.bscale != 1.0:
+                bias = bias * self.bscale
+
+        # Fast pass for linear activation.
+        if self.activation_type == 'linear' and bias is not None:
+            x = torch.addmm(bias.unsqueeze(0), x, weight.t())
+        else:
+            x = x.matmul(weight.t())
+            x = bias_act.bias_act(x, bias, act=self.activation_type, impl=impl)
+
+        assert x.dtype == dtype
+        return x
+
+# pylint: enable=missing-function-docstring

models/stylegan3_generator.py

@@ -0,0 +1,1332 @@
+# python3.7
+"""Contains the implementation of generator described in StyleGAN3.
+
+Compared to that of StyleGAN2, the generator in StyleGAN3 controls the frequency
+flow along with the convolutional layers growing.
+
+Paper: https://arxiv.org/pdf/2106.12423.pdf
+
+Official implementation: https://github.com/NVlabs/stylegan3
+"""
+
+import numpy as np
+import scipy.signal
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from third_party.stylegan3_official_ops import bias_act
+from third_party.stylegan3_official_ops import filtered_lrelu
+from third_party.stylegan3_official_ops import conv2d_gradfix
+from .utils.ops import all_gather
+
+__all__ = ['StyleGAN3Generator']
+
+# Resolutions allowed.
+_RESOLUTIONS_ALLOWED = [8, 16, 32, 64, 128, 256, 512, 1024]
+
+# pylint: disable=missing-function-docstring
+
+class StyleGAN3Generator(nn.Module):
+    """Defines the generator network in StyleGAN3.
+
+    NOTE: The synthesized images are with `RGB` channel order and pixel range
+    [-1, 1].
+
+    Settings for the mapping network:
+
+    (1) z_dim: Dimension of the input latent space, Z. (default: 512)
+    (2) w_dim: Dimension of the output latent space, W. (default: 512)
+    (3) repeat_w: Repeat w-code for different layers. (default: True)
+    (4) normalize_z: Whether to normalize the z-code. (default: True)
+    (5) mapping_layers: Number of layers of the mapping network. (default: 2)
+    (6) mapping_fmaps: Number of hidden channels of the mapping network.
+        (default: 512)
+    (7) mapping_lr_mul: Learning rate multiplier for the mapping network.
+        (default: 0.01)
+
+    Settings for conditional generation:
+
+    (1) label_dim: Dimension of the additional label for conditional generation.
+        In one-hot conditioning case, it is equal to the number of classes. If
+        set to 0, conditioning training will be disabled. (default: 0)
+    (2) embedding_dim: Dimension of the embedding space, if needed.
+        (default: 512)
+    (3) embedding_bias: Whether to add bias to embedding learning.
+        (default: True)
+    (4) embedding_lr_mul: Learning rate multiplier for the embedding learning.
+        (default: 1.0)
+    (5) normalize_embedding: Whether to normalize the embedding. (default: True)
+    (6) normalize_embedding_latent: Whether to normalize the embedding together
+        with the latent. (default: False)
+
+    Settings for the synthesis network:
+
+    (1) resolution: The resolution of the output image. (default: -1)
+    (2) image_channels: Number of channels of the output image. (default: 3)
+    (3) final_tanh: Whether to use `tanh` to control the final pixel range.
+        (default: False)
+    (4) output_scale: Factor to scaling the output image. (default: 0.25)
+    (5) num_layers: Number of synthesis layers, excluding the first positional
+        encoding layer and the last ToRGB layer. (default: 14)
+    (6) num_critical: Number of synthesis layers with critical sampling. These
+        layers are always set as top (with highest resolution) ones.
+    (7) fmaps_base: Factor to control number of feature maps for each layer.
+         (default: 32 << 10)
+    (8) fmaps_max: Maximum number of feature maps in each layer. (default: 512)
+    (9) kernel_size: Size of convolutional kernels. (default: 1)
+    (10) conv_clamp: A threshold to clamp the output of convolution layers to
+         avoid overflow under FP16 training. (default: None)
+    (11) first_cutoff: Cutoff frequency of the first layer. (default: 2)
+    (12) first_stopband: Stopband of the first layer. (default: 2 ** 2.1)
+    (13) last_stopband_rel: Stopband of the last layer, relative to the last
+         cutoff, which is `resolution / 2`. Concretely, `last_stopband` will be
+         equal to `resolution / 2 * last_stopband_rel`. (default: 2 ** 0.3)
+    (14) margin_size: Size of margin for each feature map. (default: 10)
+    (15) filter_size: Size of filter for upsampling and downsampling around the
+         activation. (default: 6)
+    (16) act_upsampling: Factor used to upsample the feature map before
+         activation for anti-aliasing. (default: 2)
+    (17) use_radial_filter: Whether to use radial filter for downsampling after
+         the activation. (default: False)
+    (18) eps: A small value to avoid divide overflow. (default: 1e-8)
+
+    Runtime settings:
+
+    (1) w_moving_decay: Decay factor for updating `w_avg`, which is used for
+        training only. Set `None` to disable. (default: 0.998)
+    (2) sync_w_avg: Synchronizing the stats of `w_avg` across replicas. If set
+        as `True`, the stats will be more accurate, yet the speed maybe a little
+        bit slower. (default: False)
+    (3) style_mixing_prob: Probability to perform style mixing as a training
+        regularization. Set `None` to disable. (default: None)
+    (4) trunc_psi: Truncation psi, set `None` to disable. (default: None)
+    (5) trunc_layers: Number of layers to perform truncation. (default: None)
+    (6) magnitude_moving_decay: Decay factor for updating `magnitude_ema` in
+        each `SynthesisLayer`, which is used for training only. Set `None` to
+        disable. (default: 0.999)
+    (7) update_ema: Whether to update `w_avg` in the `MappingNetwork` and
+        `magnitude_ema` in each `SynthesisLayer`. This field only takes effect
+        in `training` model. (default: False)
+    (8) fp16_res: Layers at resolution higher than (or equal to) this field will
+        use `float16` precision for computation. This is merely used for
+        acceleration. If set as `None`, all layers will use `float32` by
+        default. (default: None)
+    (9) impl: Implementation mode of some particular ops, e.g., `filtering`,
+        `bias_act`, etc. `cuda` means using the official CUDA implementation
+        from StyleGAN3, while `ref` means using the native PyTorch ops.
+        (default: `cuda`)
+    """
+
+    def __init__(self,
+                 # Settings for mapping network.
+                 z_dim=512,
+                 w_dim=512,
+                 repeat_w=True,
+                 normalize_z=True,
+                 mapping_layers=2,
+                 mapping_fmaps=512,
+                 mapping_lr_mul=0.01,
+                 # Settings for conditional generation.
+                 label_dim=0,
+                 embedding_dim=512,
+                 embedding_bias=True,
+                 embedding_lr_mul=1.0,
+                 normalize_embedding=True,
+                 normalize_embedding_latent=False,
+                 # Settings for synthesis network.
+                 resolution=-1,
+                 image_channels=3,
+                 final_tanh=False,
+                 output_scale=0.25,
+                 num_layers=14,
+                 num_critical=2,
+                 fmaps_base=32 << 10,
+                 fmaps_max=512,
+                 kernel_size=1,
+                 conv_clamp=256,
+                 first_cutoff=2,
+                 first_stopband=2 ** 2.1,
+                 last_stopband_rel=2 ** 0.3,
+                 margin_size=10,
+                 filter_size=6,
+                 act_upsampling=2,
+                 use_radial_filter=False,
+                 eps=1e-8):
+        """Initializes with basic settings."""
+        super().__init__()
+
+        if resolution not in _RESOLUTIONS_ALLOWED:
+            raise ValueError(f'Invalid resolution: `{resolution}`!\n'
+                             f'Resolutions allowed: {_RESOLUTIONS_ALLOWED}.')
+
+        self.z_dim = z_dim
+        self.w_dim = w_dim
+        self.repeat_w = repeat_w
+        self.normalize_z = normalize_z
+        self.mapping_layers = mapping_layers
+        self.mapping_fmaps = mapping_fmaps
+        self.mapping_lr_mul = mapping_lr_mul
+
+        self.label_dim = label_dim
+        self.embedding_dim = embedding_dim
+        self.embedding_bias = embedding_bias
+        self.embedding_lr_mul = embedding_lr_mul
+        self.normalize_embedding = normalize_embedding
+        self.normalize_embedding_latent = normalize_embedding_latent
+
+        self.resolution = resolution
+        self.image_channels = image_channels
+        self.final_tanh = final_tanh
+        self.output_scale = output_scale
+        self.num_layers = num_layers + 2  # Including InputLayer and ToRGBLayer.
+        self.num_critical = num_critical
+        self.fmaps_base = fmaps_base
+        self.fmaps_max = fmaps_max
+        self.kernel_size = kernel_size
+        self.conv_clamp = conv_clamp
+        self.first_cutoff = first_cutoff
+        self.first_stopband = first_stopband
+        self.last_stopband_rel = last_stopband_rel
+        self.margin_size = margin_size
+        self.filter_size = filter_size
+        self.act_upsampling = act_upsampling
+        self.use_radial_filter = use_radial_filter
+        self.eps = eps
+
+        # Dimension of latent space, which is convenient for sampling.
+        self.latent_dim = (z_dim,)
+
+        self.mapping = MappingNetwork(
+            input_dim=z_dim,
+            output_dim=w_dim,
+            num_outputs=self.num_layers,
+            repeat_output=repeat_w,
+            normalize_input=normalize_z,
+            num_layers=mapping_layers,
+            hidden_dim=mapping_fmaps,
+            lr_mul=mapping_lr_mul,
+            label_dim=label_dim,
+            embedding_dim=embedding_dim,
+            embedding_bias=embedding_bias,
+            embedding_lr_mul=embedding_lr_mul,
+            normalize_embedding=normalize_embedding,
+            normalize_embedding_latent=normalize_embedding_latent,
+            eps=eps)
+
+        # This is used for truncation trick.
+        if self.repeat_w:
+            self.register_buffer('w_avg', torch.zeros(w_dim))
+        else:
+            self.register_buffer('w_avg', torch.zeros(self.num_layers * w_dim))
+
+        self.synthesis = SynthesisNetwork(resolution=resolution,
+                                          w_dim=w_dim,
+                                          image_channels=image_channels,
+                                          final_tanh=final_tanh,
+                                          output_scale=output_scale,
+                                          num_layers=num_layers,
+                                          num_critical=num_critical,
+                                          fmaps_base=fmaps_base,
+                                          fmaps_max=fmaps_max,
+                                          kernel_size=kernel_size,
+                                          conv_clamp=conv_clamp,
+                                          first_cutoff=first_cutoff,
+                                          first_stopband=first_stopband,
+                                          last_stopband_rel=last_stopband_rel,
+                                          margin_size=margin_size,
+                                          filter_size=filter_size,
+                                          act_upsampling=act_upsampling,
+                                          use_radial_filter=use_radial_filter,
+                                          eps=eps)
+
+        self.var_mapping = {'w_avg': 'mapping.w_avg'}
+        for key, val in self.mapping.var_mapping.items():
+            self.var_mapping[f'mapping.{key}'] = f'mapping.{val}'
+        for key, val in self.synthesis.var_mapping.items():
+            self.var_mapping[f'synthesis.{key}'] = f'synthesis.{val}'
+
+    def set_space_of_latent(self, space_of_latent):
+        """Sets the space to which the latent code belong.
+
+        See `SynthesisNetwork` for more details.
+        """
+        self.synthesis.set_space_of_latent(space_of_latent)
+
+    def forward(self,
+                z,
+                label=None,
+                w_moving_decay=0.998,
+                sync_w_avg=False,
+                style_mixing_prob=None,
+                trunc_psi=None,
+                trunc_layers=None,
+                magnitude_moving_decay=0.999,
+                update_ema=False,
+                fp16_res=None,
+                impl='cuda'):
+        """Connects mapping network and synthesis network.
+
+        This forward function will also update the average `w_code`, perform
+        style mixing as a training regularizer, and do truncation trick, which
+        is specially designed for inference.
+
+        Concretely, the truncation trick acts as follows:
+
+        For layers in range [0, truncation_layers), the truncated w-code is
+        computed as
+
+        w_new = w_avg + (w - w_avg) * truncation_psi
+
+        To disable truncation, please set
+
+        (1) truncation_psi = 1.0 (None) OR
+        (2) truncation_layers = 0 (None)
+        """
+
+        mapping_results = self.mapping(z, label, impl=impl)
+
+        w = mapping_results['w']
+        if self.training and update_ema and w_moving_decay is not None:
+            if sync_w_avg:
+                batch_w_avg = all_gather(w.detach()).mean(dim=0)
+            else:
+                batch_w_avg = w.detach().mean(dim=0)
+            self.w_avg.copy_(batch_w_avg.lerp(self.w_avg, w_moving_decay))
+
+        wp = mapping_results.pop('wp')
+        if self.training and style_mixing_prob is not None:
+            if np.random.uniform() < style_mixing_prob:
+                new_z = torch.randn_like(z)
+                new_wp = self.mapping(new_z, label, impl=impl)['wp']
+                mixing_cutoff = np.random.randint(1, self.num_layers)
+                wp[:, mixing_cutoff:] = new_wp[:, mixing_cutoff:]
+
+        if not self.training:
+            trunc_psi = 1.0 if trunc_psi is None else trunc_psi
+            trunc_layers = 0 if trunc_layers is None else trunc_layers
+            if trunc_psi < 1.0 and trunc_layers > 0:
+                w_avg = self.w_avg.reshape(1, -1, self.w_dim)[:, :trunc_layers]
+                wp[:, :trunc_layers] = w_avg.lerp(
+                    wp[:, :trunc_layers], trunc_psi)
+
+        synthesis_results = self.synthesis(
+            wp,
+            magnitude_moving_decay=magnitude_moving_decay,
+            update_ema=update_ema,
+            fp16_res=fp16_res,
+            impl=impl)
+
+        return {**mapping_results, **synthesis_results}
+
+
+class MappingNetwork(nn.Module):
+    """Implements the latent space mapping network.
+
+    Basically, this network executes several dense layers in sequence, and the
+    label embedding if needed.
+    """
+
+    def __init__(self,
+                 input_dim,
+                 output_dim,
+                 num_outputs,
+                 repeat_output,
+                 normalize_input,
+                 num_layers,
+                 hidden_dim,
+                 lr_mul,
+                 label_dim,
+                 embedding_dim,
+                 embedding_bias,
+                 embedding_lr_mul,
+                 normalize_embedding,
+                 normalize_embedding_latent,
+                 eps):
+        super().__init__()
+
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+        self.num_outputs = num_outputs
+        self.repeat_output = repeat_output
+        self.normalize_input = normalize_input
+        self.num_layers = num_layers
+        self.hidden_dim = hidden_dim
+        self.lr_mul = lr_mul
+        self.label_dim = label_dim
+        self.embedding_dim = embedding_dim
+        self.embedding_bias = embedding_bias
+        self.embedding_lr_mul = embedding_lr_mul
+        self.normalize_embedding = normalize_embedding
+        self.normalize_embedding_latent = normalize_embedding_latent
+        self.eps = eps
+
+        self.var_mapping = {}
+
+        self.norm = PixelNormLayer(dim=1, eps=eps)
+
+        if self.label_dim > 0:
+            input_dim = input_dim + embedding_dim
+            self.embedding = DenseLayer(in_channels=label_dim,
+                                        out_channels=embedding_dim,
+                                        init_weight_std=1.0,
+                                        add_bias=embedding_bias,
+                                        init_bias=0.0,
+                                        lr_mul=embedding_lr_mul,
+                                        activation_type='linear')
+            self.var_mapping['embedding.weight'] = 'embed.weight'
+            if self.embedding_bias:
+                self.var_mapping['embedding.bias'] = 'embed.bias'
+
+        if num_outputs is not None and not repeat_output:
+            output_dim = output_dim * num_outputs
+        for i in range(num_layers):
+            in_channels = (input_dim if i == 0 else hidden_dim)
+            out_channels = (output_dim if i == (num_layers - 1) else hidden_dim)
+            self.add_module(f'dense{i}',
+                            DenseLayer(in_channels=in_channels,
+                                       out_channels=out_channels,
+                                       init_weight_std=1.0,
+                                       add_bias=True,
+                                       init_bias=0.0,
+                                       lr_mul=lr_mul,
+                                       activation_type='lrelu'))
+            self.var_mapping[f'dense{i}.weight'] = f'fc{i}.weight'
+            self.var_mapping[f'dense{i}.bias'] = f'fc{i}.bias'
+
+    def forward(self, z, label=None, impl='cuda'):
+        if z.ndim != 2 or z.shape[1] != self.input_dim:
+            raise ValueError(f'Input latent code should be with shape '
+                             f'[batch_size, input_dim], where '
+                             f'`input_dim` equals to {self.input_dim}!\n'
+                             f'But `{z.shape}` is received!')
+        if self.normalize_input:
+            z = self.norm(z)
+
+        if self.label_dim > 0:
+            if label is None:
+                raise ValueError(f'Model requires an additional label '
+                                 f'(with dimension {self.label_dim}) as input, '
+                                 f'but no label is received!')
+            if label.ndim != 2 or label.shape != (z.shape[0], self.label_dim):
+                raise ValueError(f'Input label should be with shape '
+                                 f'[batch_size, label_dim], where '
+                                 f'`batch_size` equals to that of '
+                                 f'latent codes ({z.shape[0]}) and '
+                                 f'`label_dim` equals to {self.label_dim}!\n'
+                                 f'But `{label.shape}` is received!')
+            label = label.to(dtype=torch.float32)
+            embedding = self.embedding(label, impl=impl)
+            if self.normalize_embedding:
+                embedding = self.norm(embedding)
+            w = torch.cat((z, embedding), dim=1)
+        else:
+            w = z
+
+        if self.label_dim > 0 and self.normalize_embedding_latent:
+            w = self.norm(w)
+
+        for i in range(self.num_layers):
+            w = getattr(self, f'dense{i}')(w, impl=impl)
+
+        wp = None
+        if self.num_outputs is not None:
+            if self.repeat_output:
+                wp = w.unsqueeze(1).repeat((1, self.num_outputs, 1))
+            else:
+                wp = w.reshape(-1, self.num_outputs, self.output_dim)
+
+        results = {
+            'z': z,
+            'label': label,
+            'w': w,
+            'wp': wp,
+        }
+        if self.label_dim > 0:
+            results['embedding'] = embedding
+        return results
+
+
+class SynthesisNetwork(nn.Module):
+    """Implements the image synthesis network.
+
+    Basically, this network executes several convolutional layers in sequence.
+    """
+
+    def __init__(self,
+                 resolution,
+                 w_dim,
+                 image_channels,
+                 final_tanh,
+                 output_scale,
+                 num_layers,
+                 num_critical,
+                 fmaps_base,
+                 fmaps_max,
+                 kernel_size,
+                 conv_clamp,
+                 first_cutoff,
+                 first_stopband,
+                 last_stopband_rel,
+                 margin_size,
+                 filter_size,
+                 act_upsampling,
+                 use_radial_filter,
+                 eps):
+        super().__init__()
+
+        self.resolution = resolution
+        self.w_dim = w_dim
+        self.image_channels = image_channels
+        self.final_tanh = final_tanh
+        self.output_scale = output_scale
+        self.num_layers = num_layers
+        self.num_critical = num_critical
+        self.fmaps_base = fmaps_base
+        self.fmaps_max = fmaps_max
+        self.kernel_size = kernel_size
+        self.conv_clamp = conv_clamp
+        self.first_cutoff = first_cutoff
+        self.first_stopband = first_stopband
+        self.last_stopband_rel = last_stopband_rel
+        self.margin_size = margin_size
+        self.filter_size = filter_size
+        self.act_upsampling = act_upsampling
+        self.use_radial_filter = use_radial_filter
+        self.eps = eps
+
+        self.var_mapping = {}
+
+        # Get layer settings.
+        last_cutoff = resolution / 2
+        last_stopband = last_cutoff * last_stopband_rel
+        layer_indices = np.arange(num_layers + 1)
+        exponents = np.minimum(layer_indices / (num_layers - num_critical), 1)
+        cutoffs = first_cutoff * (last_cutoff / first_cutoff) ** exponents
+        stopbands = (
+            first_stopband * (last_stopband / first_stopband) ** exponents)
+        sampling_rates = np.exp2(np.ceil(np.log2(
+            np.minimum(stopbands * 2, self.resolution))))
+        sampling_rates = np.int64(sampling_rates)
+        half_widths = np.maximum(stopbands, sampling_rates / 2) - cutoffs
+        sizes = sampling_rates + margin_size * 2
+        sizes[-2:] = resolution
+        sizes = np.int64(sizes)
+        channels = np.rint(np.minimum((fmaps_base / 2) / cutoffs, fmaps_max))
+        channels[-1] = image_channels
+        channels = np.int64(channels)
+
+        self.cutoffs = cutoffs
+        self.stopbands = stopbands
+        self.sampling_rates = sampling_rates
+        self.half_widths = half_widths
+        self.sizes = sizes
+        self.channels = channels
+
+        # Input layer, with positional encoding.
+        self.early_layer = InputLayer(w_dim=w_dim,
+                                      channels=channels[0],
+                                      size=sizes[0],
+                                      sampling_rate=sampling_rates[0],
+                                      cutoff=cutoffs[0])
+        self.var_mapping['early_layer.weight'] = 'input.weight'
+        self.var_mapping['early_layer.affine.weight'] = 'input.affine.weight'
+        self.var_mapping['early_layer.affine.bias'] = 'input.affine.bias'
+        self.var_mapping['early_layer.transform'] = 'input.transform'
+        self.var_mapping['early_layer.frequency'] = 'input.freqs'
+        self.var_mapping['early_layer.phase'] = 'input.phases'
+
+        # Convolutional layers.
+        for idx in range(num_layers + 1):
+            # Position related settings.
+            if idx < num_layers:
+                kernel_size = self.kernel_size
+                demodulate = True
+                act_upsampling = self.act_upsampling
+            else:  # ToRGB layer.
+                kernel_size = 1
+                demodulate = False
+                act_upsampling = 1
+            if idx < num_layers - num_critical:  # Non-critical sampling.
+                use_radial_filter = self.use_radial_filter
+            else:  # Critical sampling.
+                use_radial_filter = False
+
+            prev_idx = max(idx - 1, 0)
+            layer_name = f'layer{idx}'
+            official_layer_name = f'L{idx}_{sizes[idx]}_{channels[idx]}'
+            self.add_module(
+                layer_name,
+                SynthesisLayer(in_channels=channels[prev_idx],
+                               out_channels=channels[idx],
+                               w_dim=w_dim,
+                               kernel_size=kernel_size,
+                               demodulate=demodulate,
+                               eps=eps,
+                               conv_clamp=conv_clamp,
+                               in_size=sizes[prev_idx],
+                               out_size=sizes[idx],
+                               in_sampling_rate=sampling_rates[prev_idx],
+                               out_sampling_rate=sampling_rates[idx],
+                               in_cutoff=cutoffs[prev_idx],
+                               out_cutoff=cutoffs[idx],
+                               in_half_width=half_widths[prev_idx],
+                               out_half_width=half_widths[idx],
+                               filter_size=filter_size,
+                               use_radial_filter=use_radial_filter,
+                               act_upsampling=act_upsampling))
+
+            self.var_mapping[f'{layer_name}.magnitude_ema'] = (
+                f'{official_layer_name}.magnitude_ema')
+            self.var_mapping[f'{layer_name}.conv.weight'] = (
+                f'{official_layer_name}.weight')
+            self.var_mapping[f'{layer_name}.conv.style.weight'] = (
+                f'{official_layer_name}.affine.weight')
+            self.var_mapping[f'{layer_name}.conv.style.bias'] = (
+                f'{official_layer_name}.affine.bias')
+            self.var_mapping[f'{layer_name}.filter.bias'] = (
+                f'{official_layer_name}.bias')
+            if idx < num_layers:  # ToRGB layer does not need filters.
+                self.var_mapping[f'{layer_name}.filter.up_filter'] = (
+                    f'{official_layer_name}.up_filter')
+                self.var_mapping[f'{layer_name}.filter.down_filter'] = (
+                    f'{official_layer_name}.down_filter')
+
+    def set_space_of_latent(self, space_of_latent):
+        """Sets the space to which the latent code belong.
+
+        This function is particularly used for choosing how to inject the latent
+        code into the convolutional layers. The original generator will take a
+        W-Space code and apply it for style modulation after an affine
+        transformation. But, sometimes, it may need to directly feed an already
+        affine-transformed code into the convolutional layer, e.g., when
+        training an encoder for GAN inversion. We term the transformed space as
+        Style Space (or Y-Space). This function is designed to tell the
+        convolutional layers how to use the input code.
+
+        Args:
+            space_of_latent: The space to which the latent code belong. Case
+                insensitive. Support `W` and `Y`.
+        """
+        space_of_latent = space_of_latent.upper()
+        for module in self.modules():
+            if isinstance(module, ModulateConvLayer):
+                setattr(module, 'space_of_latent', space_of_latent)
+
+    def forward(self,
+                wp,
+                magnitude_moving_decay=0.999,
+                update_ema=False,
+                fp16_res=None,
+                impl='cuda'):
+        results = {'wp': wp}
+
+        x = self.early_layer(wp[:, 0])
+        for idx, sampling_rate in enumerate(self.sampling_rates):
+            if fp16_res is not None and sampling_rate >= fp16_res:
+                x = x.to(torch.float16)
+            layer = getattr(self, f'layer{idx}')
+            x, style = layer(x, wp[:, idx + 1],
+                             magnitude_moving_decay=magnitude_moving_decay,
+                             update_ema=update_ema,
+                             impl=impl)
+            results[f'style{idx}'] = style
+
+        if self.output_scale != 1:
+            x = x * self.output_scale
+        x = x.to(torch.float32)
+        if self.final_tanh:
+            x = torch.tanh(x)
+        results['image'] = x
+        return results
+
+
+class PixelNormLayer(nn.Module):
+    """Implements pixel-wise feature vector normalization layer."""
+
+    def __init__(self, dim, eps):
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+
+    def extra_repr(self):
+        return f'dim={self.dim}, epsilon={self.eps}'
+
+    def forward(self, x):
+        scale = (x.square().mean(dim=self.dim, keepdim=True) + self.eps).rsqrt()
+        return x * scale
+
+
+class InputLayer(nn.Module):
+    """Implements the input layer with positional encoding.
+
+    Basically, this block outputs a feature map with shape
+    `(channels, size, size)` based on the coordinate information.
+    `sampling_rate` and `cutoff` are used to control the coordinate range and
+    strength respectively.
+
+    For a low-pass filter, `cutoff` is the same as the `bandwidth`.
+    The initial frequency of the starting feature map is controlled by the
+    positional encoding `sin(2 * pi * x)`, where
+    `x = trans(coord) * frequency + phase`. We would like to introduce rich
+    information (i.e. frequencies), but keep all frequencies lower than
+    stopband, which is `sampling_rate / 2`.
+
+    Besides, this layer also supports learning a transformation from the latent
+    code w, and providing a customized transformation for inference. Please
+    use the buffer `transform`.
+
+    NOTE: `size` is different from `sampling_rate`. `sampling_rate` is the
+    actual size of the current stage, which determines the maximum frequency
+    that the feature maps can hold. `size` is the actual height and width of the
+    current feature map, including the extended border.
+    """
+
+    def __init__(self, w_dim, channels, size, sampling_rate, cutoff):
+        super().__init__()
+
+        self.w_dim = w_dim
+        self.channels = channels
+        self.size = size
+        self.sampling_rate = sampling_rate
+        self.cutoff = cutoff
+
+        # Coordinate of the entire feature map, with resolution (size, size).
+        # The coordinate range for the central (sampling_rate, sampling_rate)
+        # region is set as (-0.0, 0.5), which extends to the remaining region.
+        theta = torch.eye(2, 3)
+        theta[0, 0] = 0.5 / sampling_rate * size
+        theta[1, 1] = 0.5 / sampling_rate * size
+        grid = F.affine_grid(theta=theta.unsqueeze(0),
+                             size=(1, 1, size, size),
+                             align_corners=False)
+        self.register_buffer('grid', grid)
+
+        # Draw random frequency from a uniform 2D disc for each channel
+        # regarding X and Y dimension. And also draw a random phase for each
+        # channel. Accordingly, each channel has three pre-defined parameters,
+        # which are X-frequency, Y-frequency, and phase.
+        frequency = torch.randn(channels, 2)
+        radius = frequency.square().sum(dim=1, keepdim=True).sqrt()
+        frequency = frequency / (radius * radius.square().exp().pow(0.25))
+        frequency = frequency * cutoff
+        self.register_buffer('frequency', frequency)
+        phase = torch.rand(channels) - 0.5
+        self.register_buffer('phase', phase)
+
+        # This layer is used to map the latent code w to transform factors,
+        # with order: cos(angle), sin(angle), transpose_x, transpose_y.
+        self.affine = DenseLayer(in_channels=w_dim,
+                                 out_channels=4,
+                                 init_weight_std=0.0,
+                                 add_bias=True,
+                                 init_bias=(1, 0, 0, 0),
+                                 lr_mul=1.0,
+                                 activation_type='linear')
+
+        # It is possible to use this buffer to customize the transform of the
+        # output synthesis.
+        self.register_buffer('transform', torch.eye(3))
+
+        # Use 1x1 conv to convert positional encoding to features.
+        self.weight = nn.Parameter(torch.randn(channels, channels))
+        self.weight_scale = 1 / np.sqrt(channels)
+
+    def extra_repr(self):
+        return (f'channels={self.channels}, '
+                f'size={self.size}, '
+                f'sampling_rate={self.sampling_rate}, '
+                f'cutoff={self.cutoff:.3f}, ')
+
+    def forward(self, w):
+        batch = w.shape[0]
+
+        # Get transformation matrix.
+        # Factor controlled by latent code.
+        transformation_factor = self.affine(w)
+        # Ensure the range of cosine and sine value (first two dimension).
+        _norm = transformation_factor[:, :2].norm(dim=1, keepdim=True)
+        transformation_factor = transformation_factor / _norm
+        # Rotation.
+        rotation = torch.eye(3, device=w.device).unsqueeze(0)
+        rotation = rotation.repeat((batch, 1, 1))
+        rotation[:, 0, 0] = transformation_factor[:, 0]
+        rotation[:, 0, 1] = -transformation_factor[:, 1]
+        rotation[:, 1, 0] = transformation_factor[:, 1]
+        rotation[:, 1, 1] = transformation_factor[:, 0]
+        # Translation.
+        translation = torch.eye(3, device=w.device).unsqueeze(0)
+        translation = translation.repeat((batch, 1, 1))
+        translation[:, 0, 2] = -transformation_factor[:, 2]
+        translation[:, 1, 2] = -transformation_factor[:, 3]
+        # Customized transformation.
+        transform = rotation @ translation @ self.transform.unsqueeze(0)
+
+        # Transform frequency and shift, which is equivalent to transforming
+        # the coordinate. For example, given a coordinate, X, we would like to
+        # first transform it with the rotation matrix, R, and the translation
+        # matrix, T, as X' = RX + T. Then, we will apply frequency, f, and
+        # phase, p, with sin(2 * pi * (fX' + p)). Natively, we have
+        # fX' + p = f(RX + T) + p = (fR)X + (fT + p)
+        frequency = self.frequency.unsqueeze(0) @ transform[:, :2, :2]  # [NC2]
+        phase = self.frequency.unsqueeze(0) @ transform[:, :2, 2:]      # [NC]
+        phase = phase.squeeze(2) + self.phase.unsqueeze(0)              # [NC]
+
+        # Positional encoding.
+        x = self.grid                                                # [NHW2]
+        x = x.unsqueeze(3)                                           # [NHW12]
+        x = x @ frequency.transpose(1, 2).unsqueeze(1).unsqueeze(2)  # [NHW1C]
+        x = x.squeeze(3)                                             # [NHWC]
+        x = x + phase.unsqueeze(1).unsqueeze(2)                      # [NHWC]
+        x = torch.sin(2 * np.pi * x)                                 # [NHWC]
+
+        # Dampen out-of-band frequency that may be introduced by the customized
+        # transform `self.transform`.
+        frequency_norm = frequency.norm(dim=2)
+        stopband = self.sampling_rate / 2
+        factor = (frequency_norm - self.cutoff) / (stopband - self.cutoff)
+        amplitude = (1 - factor).clamp(0, 1)         # [NC]
+        x = x * amplitude.unsqueeze(1).unsqueeze(2)  # [NHWC]
+
+        # Project positional encoding to features.
+        weight = self.weight * self.weight_scale
+        x = x @ weight.t()
+
+        return x.permute(0, 3, 1, 2).contiguous()
+
+
+class SynthesisLayer(nn.Module):
+    """Implements the synthesis layer.
+
+    Each synthesis layer (including ToRGB layer) consists of a
+    `ModulateConvLayer` and a `FilteringActLayer`. Besides, this layer will
+    trace the magnitude (norm) of the input feature map, and update the
+    statistic with `magnitude_moving_decay`.
+    """
+
+    def __init__(self,
+                 # Settings for modulated convolution.
+                 in_channels,
+                 out_channels,
+                 w_dim,
+                 kernel_size,
+                 demodulate,
+                 eps,
+                 conv_clamp,
+                 # Settings for filtering activation.
+                 in_size,
+                 out_size,
+                 in_sampling_rate,
+                 out_sampling_rate,
+                 in_cutoff,
+                 out_cutoff,
+                 in_half_width,
+                 out_half_width,
+                 filter_size,
+                 use_radial_filter,
+                 act_upsampling):
+        """Initializes with layer settings.
+
+        Args:
+            in_channels: Number of channels of the input tensor.
+            out_channels: Number of channels of the output tensor.
+            w_dim: Dimension of W space for style modulation.
+            kernel_size: Size of the convolutional kernels.
+            demodulate: Whether to perform style demodulation.
+            eps: A small value to avoid divide overflow.
+            conv_clamp: A threshold to clamp the output of convolution layers to
+                avoid overflow under FP16 training.
+            in_size: Size of the input feature map, i.e., height and width.
+            out_size: Size of the output feature map, i.e., height and width.
+            in_sampling_rate: Sampling rate of the input feature map. Different
+                from `in_size` that includes extended border, this field
+                controls the actual maximum frequency that can be represented
+                by the feature map.
+            out_sampling_rate: Sampling rate of the output feature map.
+            in_cutoff: Cutoff frequency of the input feature map.
+            out_cutoff: Cutoff frequency of the output feature map.
+            in_half_width: Half-width of the transition band of the input
+                feature map.
+            out_half_width: Half-width of the transition band of the output
+                feature map.
+            filter_size: Size of the filter used in this layer.
+            use_radial_filter: Whether to use radial filter.
+            act_upsampling: Upsampling factor used before the activation.
+                `1` means do not wrap upsampling and downsampling around the
+                activation.
+        """
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.w_dim = w_dim
+        self.kernel_size = kernel_size
+        self.demodulate = demodulate
+        self.eps = eps
+        self.conv_clamp = conv_clamp
+
+        self.in_size = in_size
+        self.out_size = out_size
+        self.in_sampling_rate = in_sampling_rate
+        self.out_sampling_rate = out_sampling_rate
+        self.in_cutoff = in_cutoff
+        self.out_cutoff = out_cutoff
+        self.in_half_width = in_half_width
+        self.out_half_width = out_half_width
+        self.filter_size = filter_size
+        self.use_radial_filter = use_radial_filter
+        self.act_upsampling = act_upsampling
+
+        self.conv = ModulateConvLayer(in_channels=in_channels,
+                                      out_channels=out_channels,
+                                      w_dim=w_dim,
+                                      kernel_size=kernel_size,
+                                      demodulate=demodulate,
+                                      eps=eps)
+        self.register_buffer('magnitude_ema', torch.ones(()))
+        self.filter = FilteringActLayer(out_channels=out_channels,
+                                        in_size=in_size,
+                                        out_size=out_size,
+                                        in_sampling_rate=in_sampling_rate,
+                                        out_sampling_rate=out_sampling_rate,
+                                        in_cutoff=in_cutoff,
+                                        out_cutoff=out_cutoff,
+                                        in_half_width=in_half_width,
+                                        out_half_width=out_half_width,
+                                        filter_size=filter_size,
+                                        use_radial_filter=use_radial_filter,
+                                        conv_padding=self.conv.padding,
+                                        act_upsampling=act_upsampling)
+
+    def extra_repr(self):
+        return f'conv_clamp={self.conv_clamp}'
+
+    def forward(self,
+                x,
+                w,
+                magnitude_moving_decay=0.999,
+                update_ema=False,
+                impl='cuda'):
+        if self.training and update_ema and magnitude_moving_decay is not None:
+            magnitude = x.detach().to(torch.float32).square().mean()
+            self.magnitude_ema.copy_(
+                magnitude.lerp(self.magnitude_ema, magnitude_moving_decay))
+
+        input_gain = self.magnitude_ema.rsqrt()
+        x, style = self.conv(x, w, gain=input_gain, impl=impl)
+        if self.act_upsampling > 1:
+            x = self.filter(x, np.sqrt(2), 0.2, self.conv_clamp, impl=impl)
+        else:
+            x = self.filter(x, 1, 1, self.conv_clamp, impl=impl)
+
+        return x, style
+
+
+class ModulateConvLayer(nn.Module):
+    """Implements the convolutional layer with style modulation.
+
+    Different from the one introduced in StyleGAN2, this layer has following
+    changes:
+
+    (1) fusing `conv` and `style modulation` into one op by default
+    (2) NOT adding a noise onto the output feature map.
+    (3) NOT activating the feature map, which is moved to `FilteringActLayer`.
+    """
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 w_dim,
+                 kernel_size,
+                 demodulate,
+                 eps):
+        """Initializes with layer settings.
+
+        Args:
+            in_channels: Number of channels of the input tensor.
+            out_channels: Number of channels of the output tensor.
+            w_dim: Dimension of W space for style modulation.
+            kernel_size: Size of the convolutional kernels.
+            demodulate: Whether to perform style demodulation.
+            eps: A small value to avoid divide overflow.
+        """
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.w_dim = w_dim
+        self.kernel_size = kernel_size
+        self.demodulate = demodulate
+        self.eps = eps
+
+        self.space_of_latent = 'W'
+
+        # Set up weight.
+        weight_shape = (out_channels, in_channels, kernel_size, kernel_size)
+        self.weight = nn.Parameter(torch.randn(*weight_shape))
+        self.wscale = 1.0 / np.sqrt(kernel_size * kernel_size * in_channels)
+        self.padding = kernel_size - 1
+
+        # Set up style.
+        self.style = DenseLayer(in_channels=w_dim,
+                                out_channels=in_channels,
+                                init_weight_std=1.0,
+                                add_bias=True,
+                                init_bias=1.0,
+                                lr_mul=1.0,
+                                activation_type='linear')
+
+    def extra_repr(self):
+        return (f'in_ch={self.in_channels}, '
+                f'out_ch={self.out_channels}, '
+                f'ksize={self.kernel_size}, '
+                f'demodulate={self.demodulate}')
+
+    def forward_style(self, w, impl='cuda'):
+        """Gets style code from the given input.
+
+        More specifically, if the input is from W-Space, it will be projected by
+        an affine transformation. If it is from the Style Space (Y-Space), no
+        operation is required.
+
+        NOTE: For codes from Y-Space, we use slicing to make sure the dimension
+        is correct, in case that the code is padded before fed into this layer.
+        """
+        space_of_latent = self.space_of_latent.upper()
+        if space_of_latent == 'W':
+            if w.ndim != 2 or w.shape[1] != self.w_dim:
+                raise ValueError(f'The input tensor should be with shape '
+                                 f'[batch_size, w_dim], where '
+                                 f'`w_dim` equals to {self.w_dim}!\n'
+                                 f'But `{w.shape}` is received!')
+            style = self.style(w, impl=impl)
+        elif space_of_latent == 'Y':
+            if w.ndim != 2 or w.shape[1] < self.in_channels:
+                raise ValueError(f'The input tensor should be with shape '
+                                 f'[batch_size, y_dim], where '
+                                 f'`y_dim` equals to {self.in_channels}!\n'
+                                 f'But `{w.shape}` is received!')
+            style = w[:, :self.in_channels]
+        else:
+            raise NotImplementedError(f'Not implemented `space_of_latent`: '
+                                      f'`{space_of_latent}`!')
+        return style
+
+    def forward(self, x, w, gain=None, impl='cuda'):
+        dtype = x.dtype
+        N, C, H, W = x.shape
+
+        # Affine on `w`.
+        style = self.forward_style(w, impl=impl)
+        if not self.demodulate:
+            _style = style * self.wscale  # Equivalent to scaling weight.
+        else:
+            _style = style
+
+        weight = self.weight
+        out_ch, in_ch, kh, kw = weight.shape
+        assert in_ch == C
+
+        # Pre-normalize inputs.
+        if self.demodulate:
+            weight = (weight *
+                      weight.square().mean(dim=(1, 2, 3), keepdim=True).rsqrt())
+            _style = _style * _style.square().mean().rsqrt()
+
+        weight = weight.unsqueeze(0)
+        weight = weight * _style.reshape(N, 1, in_ch, 1, 1)  # modulation
+        if self.demodulate:
+            decoef = (weight.square().sum(dim=(2, 3, 4)) + self.eps).rsqrt()
+            weight = weight * decoef.reshape(N, out_ch, 1, 1, 1)  # demodulation
+
+        if gain is not None:
+            gain = gain.expand(N, in_ch)
+            weight = weight * gain.reshape(N, 1, in_ch, 1, 1)
+
+        # Fuse `conv` and `style modulation` as one op, using group convolution.
+        x = x.reshape(1, N * in_ch, H, W)
+        w = weight.reshape(N * out_ch, in_ch, kh, kw).to(dtype)
+        x = conv2d_gradfix.conv2d(
+            x, w, padding=self.padding, groups=N, impl=impl)
+        x = x.reshape(N, out_ch, x.shape[2], x.shape[3])
+
+        assert x.dtype == dtype
+        assert style.dtype == torch.float32
+        return x, style
+
+
+class FilteringActLayer(nn.Module):
+    """Implements the activation, wrapped with upsampling and downsampling.
+
+    Basically, this layer executes the following operations in order:
+
+    (1) Apply bias.
+    (2) Upsample the feature map to increase sampling rate.
+    (3) Apply non-linearity as activation.
+    (4) Downsample the feature map to target size.
+
+    This layer is mostly borrowed from the official implementation:
+
+    https://github.com/NVlabs/stylegan3/blob/main/training/networks_stylegan3.py
+    """
+
+    def __init__(self,
+                 out_channels,
+                 in_size,
+                 out_size,
+                 in_sampling_rate,
+                 out_sampling_rate,
+                 in_cutoff,
+                 out_cutoff,
+                 in_half_width,
+                 out_half_width,
+                 filter_size,
+                 use_radial_filter,
+                 conv_padding,
+                 act_upsampling):
+        """Initializes with layer settings.
+
+        Args:
+            out_channels: Number of output channels, which is used for `bias`.
+            in_size: Size of the input feature map, i.e., height and width.
+            out_size: Size of the output feature map, i.e., height and width.
+            in_sampling_rate: Sampling rate of the input feature map. Different
+                from `in_size` that includes extended border, this field
+                controls the actual maximum frequency that can be represented
+                by the feature map.
+            out_sampling_rate: Sampling rate of the output feature map.
+            in_cutoff: Cutoff frequency of the input feature map.
+            out_cutoff: Cutoff frequency of the output feature map.
+            in_half_width: Half-width of the transition band of the input
+                feature map.
+            out_half_width: Half-width of the transition band of the output
+                feature map.
+            filter_size: Size of the filter used in this layer.
+            use_radial_filter: Whether to use radial filter.
+            conv_padding: The padding used in the previous convolutional layer.
+            act_upsampling: Upsampling factor used before the activation.
+                `1` means do not wrap upsampling and downsampling around the
+                activation.
+        """
+        super().__init__()
+
+        self.out_channels = out_channels
+        self.in_size = in_size
+        self.out_size = out_size
+        self.in_sampling_rate = in_sampling_rate
+        self.out_sampling_rate = out_sampling_rate
+        self.in_cutoff = in_cutoff
+        self.out_cutoff = out_cutoff
+        self.in_half_width = in_half_width
+        self.out_half_width = out_half_width
+        self.filter_size = filter_size
+        self.use_radial_filter = use_radial_filter
+        self.conv_padding = conv_padding
+        self.act_upsampling = act_upsampling
+
+        # Define bias.
+        self.bias = nn.Parameter(torch.zeros(out_channels))
+
+        # This sampling rate describes the upsampled feature map before
+        # activation.
+        temp_sampling_rate = max(in_sampling_rate, out_sampling_rate)
+        temp_sampling_rate = temp_sampling_rate * act_upsampling
+
+        # Design upsampling filter.
+        up_factor = int(np.rint(temp_sampling_rate / in_sampling_rate))
+        assert in_sampling_rate * up_factor == temp_sampling_rate
+        if up_factor > 1:
+            self.up_factor = up_factor
+            self.up_taps = filter_size * up_factor
+        else:
+            self.up_factor = 1
+            self.up_taps = 1  # No filtering.
+        self.register_buffer(
+            'up_filter',
+            self.design_lowpass_filter(numtaps=self.up_taps,
+                                       cutoff=in_cutoff,
+                                       width=in_half_width * 2,
+                                       fs=temp_sampling_rate,
+                                       radial=False))
+
+        # Design downsampling filter.
+        down_factor = int(np.rint(temp_sampling_rate / out_sampling_rate))
+        assert out_sampling_rate * down_factor == temp_sampling_rate
+        if down_factor > 1:
+            self.down_factor = down_factor
+            self.down_taps = filter_size * down_factor
+        else:
+            self.down_factor = 1
+            self.down_taps = 1  # No filtering.
+        self.register_buffer(
+            'down_filter',
+            self.design_lowpass_filter(numtaps=self.down_taps,
+                                       cutoff=out_cutoff,
+                                       width=out_half_width * 2,
+                                       fs=temp_sampling_rate,
+                                       radial=use_radial_filter))
+
+        # Compute padding.
+        # Desired output size before downsampling.
+        pad_total = (out_size - 1) * self.down_factor + 1
+        # Input size after upsampling.
+        pad_total = pad_total - (in_size + conv_padding) * self.up_factor
+        # Size reduction caused by the filters.
+        pad_total = pad_total + self.up_taps + self.down_taps - 2
+        # Shift sample locations according to the symmetric interpretation.
+        pad_lo = (pad_total + self.up_factor) // 2
+        pad_hi = pad_total - pad_lo
+        self.padding = list(map(int, (pad_lo, pad_hi, pad_lo, pad_hi)))
+
+    def extra_repr(self):
+        return (f'in_size={self.in_size}, '
+                f'out_size={self.out_size}, '
+                f'in_srate={self.in_sampling_rate}, '
+                f'out_srate={self.out_sampling_rate}, '
+                f'in_cutoff={self.in_cutoff:.3f}, '
+                f'out_cutoff={self.out_cutoff:.3f}, '
+                f'in_half_width={self.in_half_width:.3f}, '
+                f'out_half_width={self.out_half_width:.3f}, '
+                f'up_factor={self.up_factor}, '
+                f'up_taps={self.up_taps}, '
+                f'down_factor={self.down_factor}, '
+                f'down_taps={self.down_taps}, '
+                f'filter_size={self.filter_size}, '
+                f'radial_filter={self.use_radial_filter}, '
+                f'conv_padding={self.conv_padding}, '
+                f'act_upsampling={self.act_upsampling}')
+
+    @staticmethod
+    def design_lowpass_filter(numtaps, cutoff, width, fs, radial=False):
+        """Designs a low-pass filter.
+
+        Args:
+            numtaps: Length of the filter (number of coefficients, i.e., the
+                filter order + 1).
+            cutoff: Cutoff frequency of the output filter.
+            width: Width of the transition region.
+            fs: Sampling frequency.
+            radial: Whether to use radially symmetric jinc-based filter.
+                (default: False)
+        """
+        if numtaps == 1:
+            return None
+
+        assert numtaps > 1
+
+        if not radial:  # Separable Kaiser low-pass filter.
+            f = scipy.signal.firwin(numtaps=numtaps,
+                                    cutoff=cutoff,
+                                    width=width,
+                                    fs=fs)
+        else:  # Radially symmetric jinc-based filter.
+            x = (np.arange(numtaps) - (numtaps - 1) / 2) / fs
+            r = np.hypot(*np.meshgrid(x, x))
+            f = scipy.special.j1(2 * cutoff * (np.pi * r)) / (np.pi * r)
+            beta = scipy.signal.kaiser_beta(
+                scipy.signal.kaiser_atten(numtaps, width / (fs / 2)))
+            w = np.kaiser(numtaps, beta)
+            f = f * np.outer(w, w)
+            f = f / np.sum(f)
+        return torch.as_tensor(f, dtype=torch.float32)
+
+    def forward(self, x, gain, slope, clamp, impl='cuda'):
+        dtype = x.dtype
+
+        x = filtered_lrelu.filtered_lrelu(x=x,
+                                          fu=self.up_filter,
+                                          fd=self.down_filter,
+                                          b=self.bias.to(dtype),
+                                          up=self.up_factor,
+                                          down=self.down_factor,
+                                          padding=self.padding,
+                                          gain=gain,
+                                          slope=slope,
+                                          clamp=clamp,
+                                          impl=impl)
+
+        assert x.dtype == dtype
+        return x
+
+
+class DenseLayer(nn.Module):
+    """Implements the dense layer."""
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 init_weight_std,
+                 add_bias,
+                 init_bias,
+                 lr_mul,
+                 activation_type):
+        """Initializes with layer settings.
+
+        Args:
+            in_channels: Number of channels of the input tensor.
+            out_channels: Number of channels of the output tensor.
+            init_weight_std: The initial standard deviation of weight.
+            add_bias: Whether to add bias onto the fully-connected result.
+            init_bias: The initial bias value before training.
+            lr_mul: Learning multiplier for both weight and bias.
+            activation_type: Type of activation.
+        """
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.init_weight_std = init_weight_std
+        self.add_bias = add_bias
+        self.init_bias = init_bias
+        self.lr_mul = lr_mul
+        self.activation_type = activation_type
+
+        weight_shape = (out_channels, in_channels)
+        self.weight = nn.Parameter(
+            torch.randn(*weight_shape) * init_weight_std / lr_mul)
+        self.wscale = lr_mul / np.sqrt(in_channels)
+
+        if add_bias:
+            init_bias = np.float32(np.float32(init_bias) / lr_mul)
+            if isinstance(init_bias, np.float32):
+                self.bias = nn.Parameter(torch.full([out_channels], init_bias))
+            else:
+                assert isinstance(init_bias, np.ndarray)
+                self.bias = nn.Parameter(torch.from_numpy(init_bias))
+            self.bscale = lr_mul
+        else:
+            self.bias = None
+
+    def extra_repr(self):
+        return (f'in_ch={self.in_channels}, '
+                f'out_ch={self.out_channels}, '
+                f'init_weight_std={self.init_weight_std}, '
+                f'bias={self.add_bias}, '
+                f'init_bias={self.init_bias}, '
+                f'lr_mul={self.lr_mul:.3f}, '
+                f'act={self.activation_type}')
+
+    def forward(self, x, impl='cuda'):
+        dtype = x.dtype
+
+        if x.ndim != 2:
+            x = x.flatten(start_dim=1)
+
+        weight = self.weight.to(dtype) * self.wscale
+        bias = None
+        if self.bias is not None:
+            bias = self.bias.to(dtype)
+            if self.bscale != 1.0:
+                bias = bias * self.bscale
+
+        # Fast pass for linear activation.
+        if self.activation_type == 'linear' and bias is not None:
+            x = torch.addmm(bias.unsqueeze(0), x, weight.t())
+        else:
+            x = x.matmul(weight.t())
+            x = bias_act.bias_act(x, bias, act=self.activation_type, impl=impl)
+
+        assert x.dtype == dtype
+        return x
+
+# pylint: enable=missing-function-docstring

models/stylegan_discriminator.py

@@ -0,0 +1,624 @@
+# python3.7
+"""Contains the implementation of discriminator described in StyleGAN.
+
+Paper: https://arxiv.org/pdf/1812.04948.pdf
+
+Official TensorFlow implementation: https://github.com/NVlabs/stylegan
+"""
+
+import numpy as np
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.cuda.amp import autocast
+
+__all__ = ['StyleGANDiscriminator']
+
+# Resolutions allowed.
+_RESOLUTIONS_ALLOWED = [8, 16, 32, 64, 128, 256, 512, 1024]
+
+# Fused-scale options allowed.
+_FUSED_SCALE_ALLOWED = [True, False, 'auto']
+
+# pylint: disable=missing-function-docstring
+
+class StyleGANDiscriminator(nn.Module):
+    """Defines the discriminator network in StyleGAN.
+
+    NOTE: The discriminator takes images with `RGB` channel order and pixel
+    range [-1, 1] as inputs.
+
+    Settings for the backbone:
+
+    (1) resolution: The resolution of the input image. (default: -1)
+    (2) init_res: Smallest resolution of the convolutional backbone.
+        (default: 4)
+    (3) image_channels: Number of channels of the input image. (default: 3)
+    (4) fused_scale:  The strategy of fusing `conv2d` and `downsample` as one
+        operator. `True` means blocks from all resolutions will fuse. `False`
+        means blocks from all resolutions will not fuse. `auto` means blocks
+        from resolutions higher than (or equal to) `fused_scale_res` will fuse.
+        (default: `auto`)
+    (5) fused_scale_res: Minimum resolution to fuse `conv2d` and `downsample`
+        as one operator. This field only takes effect if `fused_scale` is set
+        as `auto`. (default: 128)
+    (6) use_wscale: Whether to use weight scaling. (default: True)
+    (7) wscale_gain: The factor to control weight scaling. (default: sqrt(2.0))
+    (8) lr_mul: Learning rate multiplier for backbone. (default: 1.0)
+    (9) mbstd_groups: Group size for the minibatch standard deviation layer.
+        `0` means disable. (default: 4)
+    (10) mbstd_channels: Number of new channels (appended to the original
+         feature map) after the minibatch standard deviation layer. (default: 1)
+    (11) fmaps_base: Factor to control number of feature maps for each layer.
+         (default: 16 << 10)
+    (12) fmaps_max: Maximum number of feature maps in each layer. (default: 512)
+    (13) filter_kernel: Kernel used for filtering (e.g., downsampling).
+         (default: (1, 2, 1))
+    (14) eps: A small value to avoid divide overflow. (default: 1e-8)
+
+    Settings for conditional model:
+
+    (1) label_dim: Dimension of the additional label for conditional generation.
+        In one-hot conditioning case, it is equal to the number of classes. If
+        set to 0, conditioning training will be disabled. (default: 0)
+
+    Runtime settings:
+
+    (1) enable_amp: Whether to enable automatic mixed precision training.
+        (default: False)
+    """
+
+    def __init__(self,
+                 # Settings for backbone.
+                 resolution=-1,
+                 init_res=4,
+                 image_channels=3,
+                 fused_scale='auto',
+                 fused_scale_res=128,
+                 use_wscale=True,
+                 wscale_gain=np.sqrt(2.0),
+                 lr_mul=1.0,
+                 mbstd_groups=4,
+                 mbstd_channels=1,
+                 fmaps_base=16 << 10,
+                 fmaps_max=512,
+                 filter_kernel=(1, 2, 1),
+                 eps=1e-8,
+                 # Settings for conditional model.
+                 label_dim=0):
+        """Initializes with basic settings.
+
+        Raises:
+            ValueError: If the `resolution` is not supported, or `fused_scale`
+                is not supported.
+        """
+        super().__init__()
+
+        if resolution not in _RESOLUTIONS_ALLOWED:
+            raise ValueError(f'Invalid resolution: `{resolution}`!\n'
+                             f'Resolutions allowed: {_RESOLUTIONS_ALLOWED}.')
+        if fused_scale not in _FUSED_SCALE_ALLOWED:
+            raise ValueError(f'Invalid fused-scale option: `{fused_scale}`!\n'
+                             f'Options allowed: {_FUSED_SCALE_ALLOWED}.')
+
+        self.init_res = init_res
+        self.init_res_log2 = int(np.log2(init_res))
+        self.resolution = resolution
+        self.final_res_log2 = int(np.log2(resolution))
+        self.image_channels = image_channels
+        self.fused_scale = fused_scale
+        self.fused_scale_res = fused_scale_res
+        self.use_wscale = use_wscale
+        self.wscale_gain = wscale_gain
+        self.lr_mul = lr_mul
+        self.mbstd_groups = mbstd_groups
+        self.mbstd_channels = mbstd_channels
+        self.fmaps_base = fmaps_base
+        self.fmaps_max = fmaps_max
+        self.filter_kernel = filter_kernel
+        self.eps = eps
+        self.label_dim = label_dim
+
+        # Level-of-details (used for progressive training).
+        self.register_buffer('lod', torch.zeros(()))
+        self.pth_to_tf_var_mapping = {'lod': 'lod'}
+
+        for res_log2 in range(self.final_res_log2, self.init_res_log2 - 1, -1):
+            res = 2 ** res_log2
+            in_channels = self.get_nf(res)
+            out_channels = self.get_nf(res // 2)
+            block_idx = self.final_res_log2 - res_log2
+
+            # Input convolution layer for each resolution.
+            self.add_module(
+                f'input{block_idx}',
+                ConvLayer(in_channels=image_channels,
+                          out_channels=in_channels,
+                          kernel_size=1,
+                          add_bias=True,
+                          scale_factor=1,
+                          fused_scale=False,
+                          filter_kernel=None,
+                          use_wscale=use_wscale,
+                          wscale_gain=wscale_gain,
+                          lr_mul=lr_mul,
+                          activation_type='lrelu'))
+            self.pth_to_tf_var_mapping[f'input{block_idx}.weight'] = (
+                f'FromRGB_lod{block_idx}/weight')
+            self.pth_to_tf_var_mapping[f'input{block_idx}.bias'] = (
+                f'FromRGB_lod{block_idx}/bias')
+
+            # Convolution block for each resolution (except the last one).
+            if res != self.init_res:
+                # First layer (kernel 3x3) without downsampling.
+                layer_name = f'layer{2 * block_idx}'
+                self.add_module(
+                    layer_name,
+                    ConvLayer(in_channels=in_channels,
+                              out_channels=in_channels,
+                              kernel_size=3,
+                              add_bias=True,
+                              scale_factor=1,
+                              fused_scale=False,
+                              filter_kernel=None,
+                              use_wscale=use_wscale,
+                              wscale_gain=wscale_gain,
+                              lr_mul=lr_mul,
+                              activation_type='lrelu'))
+                self.pth_to_tf_var_mapping[f'{layer_name}.weight'] = (
+                    f'{res}x{res}/Conv0/weight')
+                self.pth_to_tf_var_mapping[f'{layer_name}.bias'] = (
+                    f'{res}x{res}/Conv0/bias')
+
+                # Second layer (kernel 3x3) with downsampling
+                layer_name = f'layer{2 * block_idx + 1}'
+                self.add_module(
+                    layer_name,
+                    ConvLayer(in_channels=in_channels,
+                              out_channels=out_channels,
+                              kernel_size=3,
+                              add_bias=True,
+                              scale_factor=2,
+                              fused_scale=(res >= fused_scale_res
+                                           if fused_scale == 'auto'
+                                           else fused_scale),
+                              filter_kernel=filter_kernel,
+                              use_wscale=use_wscale,
+                              wscale_gain=wscale_gain,
+                              lr_mul=lr_mul,
+                              activation_type='lrelu'))
+                self.pth_to_tf_var_mapping[f'{layer_name}.weight'] = (
+                    f'{res}x{res}/Conv1_down/weight')
+                self.pth_to_tf_var_mapping[f'{layer_name}.bias'] = (
+                    f'{res}x{res}/Conv1_down/bias')
+
+            # Convolution block for last resolution.
+            else:
+                self.mbstd = MiniBatchSTDLayer(groups=mbstd_groups,
+                                               new_channels=mbstd_channels,
+                                               eps=eps)
+
+                # First layer (kernel 3x3) without downsampling.
+                layer_name = f'layer{2 * block_idx}'
+                self.add_module(
+                    layer_name,
+                    ConvLayer(in_channels=in_channels + mbstd_channels,
+                              out_channels=in_channels,
+                              kernel_size=3,
+                              add_bias=True,
+                              scale_factor=1,
+                              fused_scale=False,
+                              filter_kernel=None,
+                              use_wscale=use_wscale,
+                              wscale_gain=wscale_gain,
+                              lr_mul=lr_mul,
+                              activation_type='lrelu'))
+                self.pth_to_tf_var_mapping[f'{layer_name}.weight'] = (
+                    f'{res}x{res}/Conv/weight')
+                self.pth_to_tf_var_mapping[f'{layer_name}.bias'] = (
+                    f'{res}x{res}/Conv/bias')
+
+                # Second layer, as a fully-connected layer.
+                layer_name = f'layer{2 * block_idx + 1}'
+                self.add_module(
+                    f'layer{2 * block_idx + 1}',
+                    DenseLayer(in_channels=in_channels * res * res,
+                               out_channels=in_channels,
+                               add_bias=True,
+                               use_wscale=use_wscale,
+                               wscale_gain=wscale_gain,
+                               lr_mul=lr_mul,
+                               activation_type='lrelu'))
+                self.pth_to_tf_var_mapping[f'{layer_name}.weight'] = (
+                    f'{res}x{res}/Dense0/weight')
+                self.pth_to_tf_var_mapping[f'{layer_name}.bias'] = (
+                    f'{res}x{res}/Dense0/bias')
+
+                # Final dense layer to output score.
+                self.output = DenseLayer(in_channels=in_channels,
+                                         out_channels=max(label_dim, 1),
+                                         add_bias=True,
+                                         use_wscale=use_wscale,
+                                         wscale_gain=1.0,
+                                         lr_mul=lr_mul,
+                                         activation_type='linear')
+                self.pth_to_tf_var_mapping['output.weight'] = (
+                    f'{res}x{res}/Dense1/weight')
+                self.pth_to_tf_var_mapping['output.bias'] = (
+                    f'{res}x{res}/Dense1/bias')
+
+    def get_nf(self, res):
+        """Gets number of feature maps according to the given resolution."""
+        return min(self.fmaps_base // res, self.fmaps_max)
+
+    def forward(self, image, label=None, lod=None, enable_amp=False):
+        expected_shape = (self.image_channels, self.resolution, self.resolution)
+        if image.ndim != 4 or image.shape[1:] != expected_shape:
+            raise ValueError(f'The input tensor should be with shape '
+                             f'[batch_size, channel, height, width], where '
+                             f'`channel` equals to {self.image_channels}, '
+                             f'`height`, `width` equal to {self.resolution}!\n'
+                             f'But `{image.shape}` is received!')
+
+        lod = self.lod.item() if lod is None else lod
+        if lod + self.init_res_log2 > self.final_res_log2:
+            raise ValueError(f'Maximum level-of-details (lod) is '
+                             f'{self.final_res_log2 - self.init_res_log2}, '
+                             f'but `{lod}` is received!')
+
+        if self.label_dim:
+            if label is None:
+                raise ValueError(f'Model requires an additional label '
+                                 f'(with dimension {self.label_dim}) as input, '
+                                 f'but no label is received!')
+            batch = image.shape[0]
+            if (label.ndim != 2 or label.shape != (batch, self.label_dim)):
+                raise ValueError(f'Input label should be with shape '
+                                 f'[batch_size, label_dim], where '
+                                 f'`batch_size` equals to {batch}, and '
+                                 f'`label_dim` equals to {self.label_dim}!\n'
+                                 f'But `{label.shape}` is received!')
+            label = label.to(dtype=torch.float32)
+
+        with autocast(enabled=enable_amp):
+            for res_log2 in range(
+                    self.final_res_log2, self.init_res_log2 - 1, -1):
+                block_idx = current_lod = self.final_res_log2 - res_log2
+                if current_lod <= lod < current_lod + 1:
+                    x = getattr(self, f'input{block_idx}')(image)
+                elif current_lod - 1 < lod < current_lod:
+                    alpha = lod - np.floor(lod)
+                    y = getattr(self, f'input{block_idx}')(image)
+                    x = y * alpha + x * (1 - alpha)
+                if lod < current_lod + 1:
+                    if res_log2 == self.init_res_log2:
+                        x = self.mbstd(x)
+                    x = getattr(self, f'layer{2 * block_idx}')(x)
+                    x = getattr(self, f'layer{2 * block_idx + 1}')(x)
+                if lod > current_lod:
+                    image = F.avg_pool2d(
+                        image, kernel_size=2, stride=2, padding=0)
+            x = self.output(x)
+
+            if self.label_dim:
+                x = (x * label).sum(dim=1, keepdim=True)
+
+        results = {
+            'score': x,
+            'label': label
+        }
+        return results
+
+
+class MiniBatchSTDLayer(nn.Module):
+    """Implements the minibatch standard deviation layer."""
+
+    def __init__(self, groups, new_channels, eps):
+        super().__init__()
+        self.groups = groups
+        self.new_channels = new_channels
+        self.eps = eps
+
+    def extra_repr(self):
+        return (f'groups={self.groups}, '
+                f'new_channels={self.new_channels}, '
+                f'epsilon={self.eps}')
+
+    def forward(self, x):
+        if self.groups <= 1 or self.new_channels < 1:
+            return x
+
+        N, C, H, W = x.shape
+        G = min(self.groups, N)  # Number of groups.
+        nC = self.new_channels  # Number of channel groups.
+        c = C // nC             # Channels per channel group.
+
+        y = x.reshape(G, -1, nC, c, H, W)  # [GnFcHW]
+        y = y - y.mean(dim=0)              # [GnFcHW]
+        y = y.square().mean(dim=0)         # [nFcHW]
+        y = (y + self.eps).sqrt()          # [nFcHW]
+        y = y.mean(dim=(2, 3, 4))          # [nF]
+        y = y.reshape(-1, nC, 1, 1)        # [nF11]
+        y = y.repeat(G, 1, H, W)           # [NFHW]
+        x = torch.cat((x, y), dim=1)       # [N(C+F)HW]
+
+        return x
+
+
+class Blur(torch.autograd.Function):
+    """Defines blur operation with customized gradient computation."""
+
+    @staticmethod
+    def forward(ctx, x, kernel):
+        assert kernel.shape[2] == 3 and kernel.shape[3] == 3
+        ctx.save_for_backward(kernel)
+        y = F.conv2d(input=x,
+                     weight=kernel,
+                     bias=None,
+                     stride=1,
+                     padding=1,
+                     groups=x.shape[1])
+        return y
+
+    @staticmethod
+    def backward(ctx, dy):
+        kernel, = ctx.saved_tensors
+        dx = BlurBackPropagation.apply(dy, kernel)
+        return dx, None, None
+
+
+class BlurBackPropagation(torch.autograd.Function):
+    """Defines the back propagation of blur operation.
+
+    NOTE: This is used to speed up the backward of gradient penalty.
+    """
+
+    @staticmethod
+    def forward(ctx, dy, kernel):
+        ctx.save_for_backward(kernel)
+        dx = F.conv2d(input=dy,
+                      weight=kernel.flip((2, 3)),
+                      bias=None,
+                      stride=1,
+                      padding=1,
+                      groups=dy.shape[1])
+        return dx
+
+    @staticmethod
+    def backward(ctx, ddx):
+        kernel, = ctx.saved_tensors
+        ddy = F.conv2d(input=ddx,
+                       weight=kernel,
+                       bias=None,
+                       stride=1,
+                       padding=1,
+                       groups=ddx.shape[1])
+        return ddy, None, None
+
+
+class ConvLayer(nn.Module):
+    """Implements the convolutional layer.
+
+    If downsampling is needed (i.e., `scale_factor = 2`), the feature map will
+    be filtered with `filter_kernel` first. If `fused_scale` is set as `True`,
+    `conv2d` and `downsample` will be fused as one operator, using stride
+    convolution.
+    """
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 add_bias,
+                 scale_factor,
+                 fused_scale,
+                 filter_kernel,
+                 use_wscale,
+                 wscale_gain,
+                 lr_mul,
+                 activation_type):
+        """Initializes with layer settings.
+
+        Args:
+            in_channels: Number of channels of the input tensor.
+            out_channels: Number of channels of the output tensor.
+            kernel_size: Size of the convolutional kernels.
+            add_bias: Whether to add bias onto the convolutional result.
+            scale_factor: Scale factor for downsampling. `1` means skip
+                downsampling.
+            fused_scale: Whether to fuse `conv2d` and `downsample` as one
+                operator, using stride convolution.
+            filter_kernel: Kernel used for filtering.
+            use_wscale: Whether to use weight scaling.
+            wscale_gain: Gain factor for weight scaling.
+            lr_mul: Learning multiplier for both weight and bias.
+            activation_type: Type of activation.
+        """
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        self.add_bias = add_bias
+        self.scale_factor = scale_factor
+        self.fused_scale = fused_scale
+        self.filter_kernel = filter_kernel
+        self.use_wscale = use_wscale
+        self.wscale_gain = wscale_gain
+        self.lr_mul = lr_mul
+        self.activation_type = activation_type
+
+        weight_shape = (out_channels, in_channels, kernel_size, kernel_size)
+        fan_in = kernel_size * kernel_size * in_channels
+        wscale = wscale_gain / np.sqrt(fan_in)
+        if use_wscale:
+            self.weight = nn.Parameter(torch.randn(*weight_shape) / lr_mul)
+            self.wscale = wscale * lr_mul
+        else:
+            self.weight = nn.Parameter(
+                torch.randn(*weight_shape) * wscale / lr_mul)
+            self.wscale = lr_mul
+
+        if add_bias:
+            self.bias = nn.Parameter(torch.zeros(out_channels))
+            self.bscale = lr_mul
+        else:
+            self.bias = None
+
+        if scale_factor > 1:
+            assert filter_kernel is not None
+            kernel = np.array(filter_kernel, dtype=np.float32).reshape(1, -1)
+            kernel = kernel.T.dot(kernel)
+            kernel = kernel / np.sum(kernel)
+            kernel = kernel[np.newaxis, np.newaxis]
+            self.register_buffer('filter', torch.from_numpy(kernel))
+
+        if scale_factor > 1 and fused_scale:  # use stride convolution.
+            self.stride = scale_factor
+        else:
+            self.stride = 1
+        self.padding = kernel_size // 2
+
+        assert activation_type in ['linear', 'relu', 'lrelu']
+
+    def extra_repr(self):
+        return (f'in_ch={self.in_channels}, '
+                f'out_ch={self.out_channels}, '
+                f'ksize={self.kernel_size}, '
+                f'wscale_gain={self.wscale_gain:.3f}, '
+                f'bias={self.add_bias}, '
+                f'lr_mul={self.lr_mul:.3f}, '
+                f'downsample={self.scale_factor}, '
+                f'fused_scale={self.fused_scale}, '
+                f'downsample_filter={self.filter_kernel}, '
+                f'act={self.activation_type}')
+
+    def forward(self, x):
+        if self.scale_factor > 1:
+            # Disable `autocast` for customized autograd function.
+            # Please check reference:
+            # https://pytorch.org/docs/stable/notes/amp_examples.html#autocast-and-custom-autograd-functions
+            with autocast(enabled=False):
+                f = self.filter.repeat(self.in_channels, 1, 1, 1)
+                x = Blur.apply(x.float(), f)  # Always use FP32.
+
+        weight = self.weight
+        if self.wscale != 1.0:
+            weight = weight * self.wscale
+        bias = None
+        if self.bias is not None:
+            bias = self.bias
+            if self.bscale != 1.0:
+                bias = bias * self.bscale
+
+        if self.scale_factor > 1 and self.fused_scale:
+            weight = F.pad(weight, (1, 1, 1, 1, 0, 0, 0, 0), 'constant', 0.0)
+            weight = (weight[:, :, 1:, 1:] + weight[:, :, :-1, 1:] +
+                      weight[:, :, 1:, :-1] + weight[:, :, :-1, :-1]) * 0.25
+        x = F.conv2d(x,
+                     weight=weight,
+                     bias=bias,
+                     stride=self.stride,
+                     padding=self.padding)
+        if self.scale_factor > 1 and not self.fused_scale:
+            down = self.scale_factor
+            x = F.avg_pool2d(x, kernel_size=down, stride=down, padding=0)
+
+        if self.activation_type == 'linear':
+            pass
+        elif self.activation_type == 'relu':
+            x = F.relu(x, inplace=True)
+        elif self.activation_type == 'lrelu':
+            x = F.leaky_relu(x, negative_slope=0.2, inplace=True)
+        else:
+            raise NotImplementedError(f'Not implemented activation type '
+                                      f'`{self.activation_type}`!')
+
+        return x
+
+
+class DenseLayer(nn.Module):
+    """Implements the dense layer."""
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 add_bias,
+                 use_wscale,
+                 wscale_gain,
+                 lr_mul,
+                 activation_type):
+        """Initializes with layer settings.
+
+        Args:
+            in_channels: Number of channels of the input tensor.
+            out_channels: Number of channels of the output tensor.
+            add_bias: Whether to add bias onto the fully-connected result.
+            use_wscale: Whether to use weight scaling.
+            wscale_gain: Gain factor for weight scaling.
+            lr_mul: Learning multiplier for both weight and bias.
+            activation_type: Type of activation.
+        """
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.add_bias = add_bias
+        self.use_wscale = use_wscale
+        self.wscale_gain = wscale_gain
+        self.lr_mul = lr_mul
+        self.activation_type = activation_type
+
+        weight_shape = (out_channels, in_channels)
+        wscale = wscale_gain / np.sqrt(in_channels)
+        if use_wscale:
+            self.weight = nn.Parameter(torch.randn(*weight_shape) / lr_mul)
+            self.wscale = wscale * lr_mul
+        else:
+            self.weight = nn.Parameter(
+                torch.randn(*weight_shape) * wscale / lr_mul)
+            self.wscale = lr_mul
+
+        if add_bias:
+            self.bias = nn.Parameter(torch.zeros(out_channels))
+            self.bscale = lr_mul
+        else:
+            self.bias = None
+
+        assert activation_type in ['linear', 'relu', 'lrelu']
+
+    def extra_repr(self):
+        return (f'in_ch={self.in_channels}, '
+                f'out_ch={self.out_channels}, '
+                f'wscale_gain={self.wscale_gain:.3f}, '
+                f'bias={self.add_bias}, '
+                f'lr_mul={self.lr_mul:.3f}, '
+                f'act={self.activation_type}')
+
+    def forward(self, x):
+        if x.ndim != 2:
+            x = x.flatten(start_dim=1)
+
+        weight = self.weight
+        if self.wscale != 1.0:
+            weight = weight * self.wscale
+        bias = None
+        if self.bias is not None:
+            bias = self.bias
+            if self.bscale != 1.0:
+                bias = bias * self.bscale
+
+        x = F.linear(x, weight=weight, bias=bias)
+
+        if self.activation_type == 'linear':
+            pass
+        elif self.activation_type == 'relu':
+            x = F.relu(x, inplace=True)
+        elif self.activation_type == 'lrelu':
+            x = F.leaky_relu(x, negative_slope=0.2, inplace=True)
+        else:
+            raise NotImplementedError(f'Not implemented activation type '
+                                      f'`{self.activation_type}`!')
+
+        return x
+
+# pylint: enable=missing-function-docstring

models/stylegan_generator.py

@@ -0,0 +1,999 @@
+# python3.7
+"""Contains the implementation of generator described in StyleGAN.
+
+Paper: https://arxiv.org/pdf/1812.04948.pdf
+
+Official TensorFlow implementation: https://github.com/NVlabs/stylegan
+"""
+
+import numpy as np
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.cuda.amp import autocast
+
+from .utils.ops import all_gather
+
+__all__ = ['StyleGANGenerator']
+
+# Resolutions allowed.
+_RESOLUTIONS_ALLOWED = [8, 16, 32, 64, 128, 256, 512, 1024]
+
+# Fused-scale options allowed.
+_FUSED_SCALE_ALLOWED = [True, False, 'auto']
+
+# pylint: disable=missing-function-docstring
+
+class StyleGANGenerator(nn.Module):
+    """Defines the generator network in StyleGAN.
+
+    NOTE: The synthesized images are with `RGB` channel order and pixel range
+    [-1, 1].
+
+    Settings for the mapping network:
+
+    (1) z_dim: Dimension of the input latent space, Z. (default: 512)
+    (2) w_dim: Dimension of the output latent space, W. (default: 512)
+    (3) repeat_w: Repeat w-code for different layers. (default: True)
+    (4) normalize_z: Whether to normalize the z-code. (default: True)
+    (5) mapping_layers: Number of layers of the mapping network. (default: 8)
+    (6) mapping_fmaps: Number of hidden channels of the mapping network.
+        (default: 512)
+    (7) mapping_use_wscale: Whether to use weight scaling for the mapping
+        network. (default: True)
+    (8) mapping_wscale_gain: The factor to control weight scaling for the
+        mapping network (default: sqrt(2.0))
+    (9) mapping_lr_mul: Learning rate multiplier for the mapping network.
+        (default: 0.01)
+
+    Settings for conditional generation:
+
+    (1) label_dim: Dimension of the additional label for conditional generation.
+        In one-hot conditioning case, it is equal to the number of classes. If
+        set to 0, conditioning training will be disabled. (default: 0)
+    (2) embedding_dim: Dimension of the embedding space, if needed.
+        (default: 512)
+
+    Settings for the synthesis network:
+
+    (1) resolution: The resolution of the output image. (default: -1)
+    (2) init_res: The initial resolution to start with convolution. (default: 4)
+    (3) image_channels: Number of channels of the output image. (default: 3)
+    (4) final_tanh: Whether to use `tanh` to control the final pixel range.
+        (default: False)
+    (5) fused_scale:  The strategy of fusing `upsample` and `conv2d` as one
+        operator. `True` means blocks from all resolutions will fuse. `False`
+        means blocks from all resolutions will not fuse. `auto` means blocks
+        from resolutions higher than (or equal to) `fused_scale_res` will fuse.
+        (default: `auto`)
+    (6) fused_scale_res: Minimum resolution to fuse `conv2d` and `downsample`
+        as one operator. This field only takes effect if `fused_scale` is set
+        as `auto`. (default: 128)
+    (7) use_wscale: Whether to use weight scaling. (default: True)
+    (8) wscale_gain: The factor to control weight scaling. (default: sqrt(2.0))
+    (9) lr_mul: Learning rate multiplier for the synthesis network.
+        (default: 1.0)
+    (10) noise_type: Type of noise added to the convolutional results at each
+         layer. (default: `spatial`)
+    (11) fmaps_base: Factor to control number of feature maps for each layer.
+         (default: 16 << 10)
+    (12) fmaps_max: Maximum number of feature maps in each layer. (default: 512)
+    (13) filter_kernel: Kernel used for filtering (e.g., downsampling).
+         (default: (1, 2, 1))
+    (14) eps: A small value to avoid divide overflow. (default: 1e-8)
+
+    Runtime settings:
+
+    (1) w_moving_decay: Decay factor for updating `w_avg`, which is used for
+        training only. Set `None` to disable. (default: None)
+    (2) sync_w_avg: Synchronizing the stats of `w_avg` across replicas. If set
+        as `True`, the stats will be more accurate, yet the speed maybe a little
+        bit slower. (default: False)
+    (3) style_mixing_prob: Probability to perform style mixing as a training
+        regularization. Set `None` to disable. (default: None)
+    (4) trunc_psi: Truncation psi, set `None` to disable. (default: None)
+    (5) trunc_layers: Number of layers to perform truncation. (default: None)
+    (6) noise_mode: Mode of the layer-wise noise. Support `none`, `random`,
+        `const`. (default: `const`)
+    (7) enable_amp: Whether to enable automatic mixed precision training.
+        (default: False)
+    """
+
+    def __init__(self,
+                 # Settings for mapping network.
+                 z_dim=512,
+                 w_dim=512,
+                 repeat_w=True,
+                 normalize_z=True,
+                 mapping_layers=8,
+                 mapping_fmaps=512,
+                 mapping_use_wscale=True,
+                 mapping_wscale_gain=np.sqrt(2.0),
+                 mapping_lr_mul=0.01,
+                 # Settings for conditional generation.
+                 label_dim=0,
+                 embedding_dim=512,
+                 # Settings for synthesis network.
+                 resolution=-1,
+                 init_res=4,
+                 image_channels=3,
+                 final_tanh=False,
+                 fused_scale='auto',
+                 fused_scale_res=128,
+                 use_wscale=True,
+                 wscale_gain=np.sqrt(2.0),
+                 lr_mul=1.0,
+                 noise_type='spatial',
+                 fmaps_base=16 << 10,
+                 fmaps_max=512,
+                 filter_kernel=(1, 2, 1),
+                 eps=1e-8):
+        """Initializes with basic settings.
+
+        Raises:
+            ValueError: If the `resolution` is not supported, or `fused_scale`
+                is not supported.
+        """
+        super().__init__()
+
+        if resolution not in _RESOLUTIONS_ALLOWED:
+            raise ValueError(f'Invalid resolution: `{resolution}`!\n'
+                             f'Resolutions allowed: {_RESOLUTIONS_ALLOWED}.')
+        if fused_scale not in _FUSED_SCALE_ALLOWED:
+            raise ValueError(f'Invalid fused-scale option: `{fused_scale}`!\n'
+                             f'Options allowed: {_FUSED_SCALE_ALLOWED}.')
+
+        self.z_dim = z_dim
+        self.w_dim = w_dim
+        self.repeat_w = repeat_w
+        self.normalize_z = normalize_z
+        self.mapping_layers = mapping_layers
+        self.mapping_fmaps = mapping_fmaps
+        self.mapping_use_wscale = mapping_use_wscale
+        self.mapping_wscale_gain = mapping_wscale_gain
+        self.mapping_lr_mul = mapping_lr_mul
+
+        self.label_dim = label_dim
+        self.embedding_dim = embedding_dim
+
+        self.resolution = resolution
+        self.init_res = init_res
+        self.image_channels = image_channels
+        self.final_tanh = final_tanh
+        self.fused_scale = fused_scale
+        self.fused_scale_res = fused_scale_res
+        self.use_wscale = use_wscale
+        self.wscale_gain = wscale_gain
+        self.lr_mul = lr_mul
+        self.noise_type = noise_type.lower()
+        self.fmaps_base = fmaps_base
+        self.fmaps_max = fmaps_max
+        self.filter_kernel = filter_kernel
+        self.eps = eps
+
+        # Dimension of latent space, which is convenient for sampling.
+        self.latent_dim = (z_dim,)
+
+        # Number of synthesis (convolutional) layers.
+        self.num_layers = int(np.log2(resolution // init_res * 2)) * 2
+
+        self.mapping = MappingNetwork(input_dim=z_dim,
+                                      output_dim=w_dim,
+                                      num_outputs=self.num_layers,
+                                      repeat_output=repeat_w,
+                                      normalize_input=normalize_z,
+                                      num_layers=mapping_layers,
+                                      hidden_dim=mapping_fmaps,
+                                      use_wscale=mapping_use_wscale,
+                                      wscale_gain=mapping_wscale_gain,
+                                      lr_mul=mapping_lr_mul,
+                                      label_dim=label_dim,
+                                      embedding_dim=embedding_dim,
+                                      eps=eps)
+
+        # This is used for truncation trick.
+        if self.repeat_w:
+            self.register_buffer('w_avg', torch.zeros(w_dim))
+        else:
+            self.register_buffer('w_avg', torch.zeros(self.num_layers * w_dim))
+
+        self.synthesis = SynthesisNetwork(resolution=resolution,
+                                          init_res=init_res,
+                                          w_dim=w_dim,
+                                          image_channels=image_channels,
+                                          final_tanh=final_tanh,
+                                          fused_scale=fused_scale,
+                                          fused_scale_res=fused_scale_res,
+                                          use_wscale=use_wscale,
+                                          wscale_gain=wscale_gain,
+                                          lr_mul=lr_mul,
+                                          noise_type=noise_type,
+                                          fmaps_base=fmaps_base,
+                                          fmaps_max=fmaps_max,
+                                          filter_kernel=filter_kernel,
+                                          eps=eps)
+
+        self.pth_to_tf_var_mapping = {'w_avg': 'dlatent_avg'}
+        for key, val in self.mapping.pth_to_tf_var_mapping.items():
+            self.pth_to_tf_var_mapping[f'mapping.{key}'] = val
+        for key, val in self.synthesis.pth_to_tf_var_mapping.items():
+            self.pth_to_tf_var_mapping[f'synthesis.{key}'] = val
+
+    def set_space_of_latent(self, space_of_latent):
+        """Sets the space to which the latent code belong.
+
+        See `SynthesisNetwork` for more details.
+        """
+        self.synthesis.set_space_of_latent(space_of_latent)
+
+    def forward(self,
+                z,
+                label=None,
+                lod=None,
+                w_moving_decay=None,
+                sync_w_avg=False,
+                style_mixing_prob=None,
+                trunc_psi=None,
+                trunc_layers=None,
+                noise_mode='const',
+                enable_amp=False):
+        mapping_results = self.mapping(z, label)
+
+        w = mapping_results['w']
+        if self.training and w_moving_decay is not None:
+            if sync_w_avg:
+                batch_w_avg = all_gather(w.detach()).mean(dim=0)
+            else:
+                batch_w_avg = w.detach().mean(dim=0)
+            self.w_avg.copy_(batch_w_avg.lerp(self.w_avg, w_moving_decay))
+
+        wp = mapping_results.pop('wp')
+        if self.training and style_mixing_prob is not None:
+            if np.random.uniform() < style_mixing_prob:
+                new_z = torch.randn_like(z)
+                new_wp = self.mapping(new_z, label)['wp']
+                lod = self.synthesis.lod.item() if lod is None else lod
+                current_layers = self.num_layers - int(lod) * 2
+                mixing_cutoff = np.random.randint(1, current_layers)
+                wp[:, mixing_cutoff:] = new_wp[:, mixing_cutoff:]
+
+        if not self.training:
+            trunc_psi = 1.0 if trunc_psi is None else trunc_psi
+            trunc_layers = 0 if trunc_layers is None else trunc_layers
+            if trunc_psi < 1.0 and trunc_layers > 0:
+                w_avg = self.w_avg.reshape(1, -1, self.w_dim)[:, :trunc_layers]
+                wp[:, :trunc_layers] = w_avg.lerp(
+                    wp[:, :trunc_layers], trunc_psi)
+
+        with autocast(enabled=enable_amp):
+            synthesis_results = self.synthesis(wp,
+                                               lod=lod,
+                                               noise_mode=noise_mode)
+
+        return {**mapping_results, **synthesis_results}
+
+
+class MappingNetwork(nn.Module):
+    """Implements the latent space mapping module.
+
+    Basically, this module executes several dense layers in sequence, and the
+    label embedding if needed.
+    """
+
+    def __init__(self,
+                 input_dim,
+                 output_dim,
+                 num_outputs,
+                 repeat_output,
+                 normalize_input,
+                 num_layers,
+                 hidden_dim,
+                 use_wscale,
+                 wscale_gain,
+                 lr_mul,
+                 label_dim,
+                 embedding_dim,
+                 eps):
+        super().__init__()
+
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+        self.num_outputs = num_outputs
+        self.repeat_output = repeat_output
+        self.normalize_input = normalize_input
+        self.num_layers = num_layers
+        self.hidden_dim = hidden_dim
+        self.use_wscale = use_wscale
+        self.wscale_gain = wscale_gain
+        self.lr_mul = lr_mul
+        self.label_dim = label_dim
+        self.embedding_dim = embedding_dim
+        self.eps = eps
+
+        self.pth_to_tf_var_mapping = {}
+
+        if normalize_input:
+            self.norm = PixelNormLayer(dim=1, eps=eps)
+
+        if self.label_dim > 0:
+            input_dim = input_dim + embedding_dim
+            self.embedding = nn.Parameter(
+                torch.randn(label_dim, embedding_dim))
+            self.pth_to_tf_var_mapping['embedding'] = 'LabelConcat/weight'
+
+        if num_outputs is not None and not repeat_output:
+            output_dim = output_dim * num_outputs
+        for i in range(num_layers):
+            in_channels = (input_dim if i == 0 else hidden_dim)
+            out_channels = (output_dim if i == (num_layers - 1) else hidden_dim)
+            self.add_module(f'dense{i}',
+                            DenseLayer(in_channels=in_channels,
+                                       out_channels=out_channels,
+                                       add_bias=True,
+                                       use_wscale=use_wscale,
+                                       wscale_gain=wscale_gain,
+                                       lr_mul=lr_mul,
+                                       activation_type='lrelu'))
+            self.pth_to_tf_var_mapping[f'dense{i}.weight'] = f'Dense{i}/weight'
+            self.pth_to_tf_var_mapping[f'dense{i}.bias'] = f'Dense{i}/bias'
+
+    def forward(self, z, label=None):
+        if z.ndim != 2 or z.shape[1] != self.input_dim:
+            raise ValueError(f'Input latent code should be with shape '
+                             f'[batch_size, input_dim], where '
+                             f'`input_dim` equals to {self.input_dim}!\n'
+                             f'But `{z.shape}` is received!')
+
+        if self.label_dim > 0:
+            if label is None:
+                raise ValueError(f'Model requires an additional label '
+                                 f'(with dimension {self.label_dim}) as input, '
+                                 f'but no label is received!')
+            if label.ndim != 2 or label.shape != (z.shape[0], self.label_dim):
+                raise ValueError(f'Input label should be with shape '
+                                 f'[batch_size, label_dim], where '
+                                 f'`batch_size` equals to that of '
+                                 f'latent codes ({z.shape[0]}) and '
+                                 f'`label_dim` equals to {self.label_dim}!\n'
+                                 f'But `{label.shape}` is received!')
+            label = label.to(dtype=torch.float32)
+            embedding = torch.matmul(label, self.embedding)
+            z = torch.cat((z, embedding), dim=1)
+
+        if self.normalize_input:
+            w = self.norm(z)
+        else:
+            w = z
+
+        for i in range(self.num_layers):
+            w = getattr(self, f'dense{i}')(w)
+
+        wp = None
+        if self.num_outputs is not None:
+            if self.repeat_output:
+                wp = w.unsqueeze(1).repeat((1, self.num_outputs, 1))
+            else:
+                wp = w.reshape(-1, self.num_outputs, self.output_dim)
+
+        results = {
+            'z': z,
+            'label': label,
+            'w': w,
+            'wp': wp,
+        }
+        if self.label_dim > 0:
+            results['embedding'] = embedding
+        return results
+
+
+class SynthesisNetwork(nn.Module):
+    """Implements the image synthesis module.
+
+    Basically, this module executes several convolutional layers in sequence.
+    """
+
+    def __init__(self,
+                 resolution,
+                 init_res,
+                 w_dim,
+                 image_channels,
+                 final_tanh,
+                 fused_scale,
+                 fused_scale_res,
+                 use_wscale,
+                 wscale_gain,
+                 lr_mul,
+                 noise_type,
+                 fmaps_base,
+                 fmaps_max,
+                 filter_kernel,
+                 eps):
+        super().__init__()
+
+        self.init_res = init_res
+        self.init_res_log2 = int(np.log2(init_res))
+        self.resolution = resolution
+        self.final_res_log2 = int(np.log2(resolution))
+        self.w_dim = w_dim
+        self.image_channels = image_channels
+        self.final_tanh = final_tanh
+        self.fused_scale = fused_scale
+        self.fused_scale_res = fused_scale_res
+        self.use_wscale = use_wscale
+        self.wscale_gain = wscale_gain
+        self.lr_mul = lr_mul
+        self.noise_type = noise_type.lower()
+        self.fmaps_base = fmaps_base
+        self.fmaps_max = fmaps_max
+        self.eps = eps
+
+        self.num_layers = (self.final_res_log2 - self.init_res_log2 + 1) * 2
+
+        # Level-of-details (used for progressive training).
+        self.register_buffer('lod', torch.zeros(()))
+        self.pth_to_tf_var_mapping = {'lod': 'lod'}
+
+        for res_log2 in range(self.init_res_log2, self.final_res_log2 + 1):
+            res = 2 ** res_log2
+            in_channels = self.get_nf(res // 2)
+            out_channels = self.get_nf(res)
+            block_idx = res_log2 - self.init_res_log2
+
+            # First layer (kernel 3x3) with upsampling
+            layer_name = f'layer{2 * block_idx}'
+            if res == self.init_res:
+                self.add_module(layer_name,
+                                ModulateConvLayer(in_channels=0,
+                                                  out_channels=out_channels,
+                                                  resolution=res,
+                                                  w_dim=w_dim,
+                                                  kernel_size=None,
+                                                  add_bias=True,
+                                                  scale_factor=None,
+                                                  fused_scale=None,
+                                                  filter_kernel=None,
+                                                  use_wscale=use_wscale,
+                                                  wscale_gain=wscale_gain,
+                                                  lr_mul=lr_mul,
+                                                  noise_type=noise_type,
+                                                  activation_type='lrelu',
+                                                  use_style=True,
+                                                  eps=eps))
+                tf_layer_name = 'Const'
+                self.pth_to_tf_var_mapping[f'{layer_name}.const'] = (
+                    f'{res}x{res}/{tf_layer_name}/const')
+            else:
+                self.add_module(
+                    layer_name,
+                    ModulateConvLayer(in_channels=in_channels,
+                                      out_channels=out_channels,
+                                      resolution=res,
+                                      w_dim=w_dim,
+                                      kernel_size=3,
+                                      add_bias=True,
+                                      scale_factor=2,
+                                      fused_scale=(res >= fused_scale_res
+                                                   if fused_scale == 'auto'
+                                                   else fused_scale),
+                                      filter_kernel=filter_kernel,
+                                      use_wscale=use_wscale,
+                                      wscale_gain=wscale_gain,
+                                      lr_mul=lr_mul,
+                                      noise_type=noise_type,
+                                      activation_type='lrelu',
+                                      use_style=True,
+                                      eps=eps))
+                tf_layer_name = 'Conv0_up'
+                self.pth_to_tf_var_mapping[f'{layer_name}.weight'] = (
+                    f'{res}x{res}/{tf_layer_name}/weight')
+            self.pth_to_tf_var_mapping[f'{layer_name}.bias'] = (
+                f'{res}x{res}/{tf_layer_name}/bias')
+            self.pth_to_tf_var_mapping[f'{layer_name}.style.weight'] = (
+                f'{res}x{res}/{tf_layer_name}/StyleMod/weight')
+            self.pth_to_tf_var_mapping[f'{layer_name}.style.bias'] = (
+                f'{res}x{res}/{tf_layer_name}/StyleMod/bias')
+            self.pth_to_tf_var_mapping[f'{layer_name}.noise_strength'] = (
+                f'{res}x{res}/{tf_layer_name}/Noise/weight')
+            self.pth_to_tf_var_mapping[f'{layer_name}.noise'] = (
+                f'noise{2 * block_idx}')
+
+            # Second layer (kernel 3x3) without upsampling.
+            layer_name = f'layer{2 * block_idx + 1}'
+            self.add_module(layer_name,
+                            ModulateConvLayer(in_channels=out_channels,
+                                              out_channels=out_channels,
+                                              resolution=res,
+                                              w_dim=w_dim,
+                                              kernel_size=3,
+                                              add_bias=True,
+                                              scale_factor=1,
+                                              fused_scale=False,
+                                              filter_kernel=None,
+                                              use_wscale=use_wscale,
+                                              wscale_gain=wscale_gain,
+                                              lr_mul=lr_mul,
+                                              noise_type=noise_type,
+                                              activation_type='lrelu',
+                                              use_style=True,
+                                              eps=eps))
+            tf_layer_name = 'Conv' if res == self.init_res else 'Conv1'
+            self.pth_to_tf_var_mapping[f'{layer_name}.weight'] = (
+                f'{res}x{res}/{tf_layer_name}/weight')
+            self.pth_to_tf_var_mapping[f'{layer_name}.bias'] = (
+                f'{res}x{res}/{tf_layer_name}/bias')
+            self.pth_to_tf_var_mapping[f'{layer_name}.style.weight'] = (
+                f'{res}x{res}/{tf_layer_name}/StyleMod/weight')
+            self.pth_to_tf_var_mapping[f'{layer_name}.style.bias'] = (
+                f'{res}x{res}/{tf_layer_name}/StyleMod/bias')
+            self.pth_to_tf_var_mapping[f'{layer_name}.noise_strength'] = (
+                f'{res}x{res}/{tf_layer_name}/Noise/weight')
+            self.pth_to_tf_var_mapping[f'{layer_name}.noise'] = (
+                f'noise{2 * block_idx + 1}')
+
+            # Output convolution layer for each resolution.
+            self.add_module(f'output{block_idx}',
+                            ModulateConvLayer(in_channels=out_channels,
+                                              out_channels=image_channels,
+                                              resolution=res,
+                                              w_dim=w_dim,
+                                              kernel_size=1,
+                                              add_bias=True,
+                                              scale_factor=1,
+                                              fused_scale=False,
+                                              filter_kernel=None,
+                                              use_wscale=use_wscale,
+                                              wscale_gain=1.0,
+                                              lr_mul=lr_mul,
+                                              noise_type='none',
+                                              activation_type='linear',
+                                              use_style=False,
+                                              eps=eps))
+            self.pth_to_tf_var_mapping[f'output{block_idx}.weight'] = (
+                f'ToRGB_lod{self.final_res_log2 - res_log2}/weight')
+            self.pth_to_tf_var_mapping[f'output{block_idx}.bias'] = (
+                f'ToRGB_lod{self.final_res_log2 - res_log2}/bias')
+
+    def get_nf(self, res):
+        """Gets number of feature maps according to the given resolution."""
+        return min(self.fmaps_base // res, self.fmaps_max)
+
+    def set_space_of_latent(self, space_of_latent):
+        """Sets the space to which the latent code belong.
+
+        This function is particularly used for choosing how to inject the latent
+        code into the convolutional layers. The original generator will take a
+        W-Space code and apply it for style modulation after an affine
+        transformation. But, sometimes, it may need to directly feed an already
+        affine-transformed code into the convolutional layer, e.g., when
+        training an encoder for GAN inversion. We term the transformed space as
+        Style Space (or Y-Space). This function is designed to tell the
+        convolutional layers how to use the input code.
+
+        Args:
+            space_of_latent: The space to which the latent code belong. Case
+                insensitive. Support `W` and `Y`.
+        """
+        space_of_latent = space_of_latent.upper()
+        for module in self.modules():
+            if isinstance(module, ModulateConvLayer) and module.use_style:
+                setattr(module, 'space_of_latent', space_of_latent)
+
+    def forward(self, wp, lod=None, noise_mode='const'):
+        lod = self.lod.item() if lod is None else lod
+        if lod + self.init_res_log2 > self.final_res_log2:
+            raise ValueError(f'Maximum level-of-details (lod) is '
+                             f'{self.final_res_log2 - self.init_res_log2}, '
+                             f'but `{lod}` is received!')
+
+        results = {'wp': wp}
+        x = None
+        for res_log2 in range(self.init_res_log2, self.final_res_log2 + 1):
+            current_lod = self.final_res_log2 - res_log2
+            block_idx = res_log2 - self.init_res_log2
+            if lod < current_lod + 1:
+                layer = getattr(self, f'layer{2 * block_idx}')
+                x, style = layer(x, wp[:, 2 * block_idx], noise_mode)
+                results[f'style{2 * block_idx}'] = style
+                layer = getattr(self, f'layer{2 * block_idx + 1}')
+                x, style = layer(x, wp[:, 2 * block_idx + 1], noise_mode)
+                results[f'style{2 * block_idx + 1}'] = style
+            if current_lod - 1 < lod <= current_lod:
+                image = getattr(self, f'output{block_idx}')(x)
+            elif current_lod < lod < current_lod + 1:
+                alpha = np.ceil(lod) - lod
+                temp = getattr(self, f'output{block_idx}')(x)
+                image = F.interpolate(image, scale_factor=2, mode='nearest')
+                image = temp * alpha + image * (1 - alpha)
+            elif lod >= current_lod + 1:
+                image = F.interpolate(image, scale_factor=2, mode='nearest')
+
+        if self.final_tanh:
+            image = torch.tanh(image)
+        results['image'] = image
+        return results
+
+
+class PixelNormLayer(nn.Module):
+    """Implements pixel-wise feature vector normalization layer."""
+
+    def __init__(self, dim, eps):
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+
+    def extra_repr(self):
+        return f'dim={self.dim}, epsilon={self.eps}'
+
+    def forward(self, x):
+        scale = (x.square().mean(dim=self.dim, keepdim=True) + self.eps).rsqrt()
+        return x * scale
+
+
+class Blur(torch.autograd.Function):
+    """Defines blur operation with customized gradient computation."""
+
+    @staticmethod
+    def forward(ctx, x, kernel):
+        assert kernel.shape[2] == 3 and kernel.shape[3] == 3
+        ctx.save_for_backward(kernel)
+        y = F.conv2d(input=x,
+                     weight=kernel,
+                     bias=None,
+                     stride=1,
+                     padding=1,
+                     groups=x.shape[1])
+        return y
+
+    @staticmethod
+    def backward(ctx, dy):
+        kernel, = ctx.saved_tensors
+        dx = F.conv2d(input=dy,
+                      weight=kernel.flip((2, 3)),
+                      bias=None,
+                      stride=1,
+                      padding=1,
+                      groups=dy.shape[1])
+        return dx, None, None
+
+
+class ModulateConvLayer(nn.Module):
+    """Implements the convolutional layer with style modulation."""
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 resolution,
+                 w_dim,
+                 kernel_size,
+                 add_bias,
+                 scale_factor,
+                 fused_scale,
+                 filter_kernel,
+                 use_wscale,
+                 wscale_gain,
+                 lr_mul,
+                 noise_type,
+                 activation_type,
+                 use_style,
+                 eps):
+        """Initializes with layer settings.
+
+        Args:
+            in_channels: Number of channels of the input tensor.
+            out_channels: Number of channels of the output tensor.
+            resolution: Resolution of the output tensor.
+            w_dim: Dimension of W space for style modulation.
+            kernel_size: Size of the convolutional kernels.
+            add_bias: Whether to add bias onto the convolutional result.
+            scale_factor: Scale factor for upsampling.
+            fused_scale: Whether to fuse `upsample` and `conv2d` as one
+                operator, using transpose convolution.
+            filter_kernel: Kernel used for filtering.
+            use_wscale: Whether to use weight scaling.
+            wscale_gain: Gain factor for weight scaling.
+            lr_mul: Learning multiplier for both weight and bias.
+            noise_type: Type of noise added to the feature map after the
+                convolution (if needed). Support `none`, `spatial` and
+                `channel`.
+            activation_type: Type of activation.
+            use_style: Whether to apply style modulation.
+            eps: A small value to avoid divide overflow.
+        """
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.resolution = resolution
+        self.w_dim = w_dim
+        self.kernel_size = kernel_size
+        self.add_bias = add_bias
+        self.scale_factor = scale_factor
+        self.fused_scale = fused_scale
+        self.filter_kernel = filter_kernel
+        self.use_wscale = use_wscale
+        self.wscale_gain = wscale_gain
+        self.lr_mul = lr_mul
+        self.noise_type = noise_type.lower()
+        self.activation_type = activation_type
+        self.use_style = use_style
+        self.eps = eps
+
+        # Set up noise.
+        if self.noise_type == 'none':
+            pass
+        elif self.noise_type == 'spatial':
+            self.register_buffer(
+                'noise', torch.randn(1, 1, resolution, resolution))
+            self.noise_strength = nn.Parameter(
+                torch.zeros(1, out_channels, 1, 1))
+        elif self.noise_type == 'channel':
+            self.register_buffer(
+                'noise', torch.randn(1, out_channels, 1, 1))
+            self.noise_strength = nn.Parameter(
+                torch.zeros(1, 1, resolution, resolution))
+        else:
+            raise NotImplementedError(f'Not implemented noise type: '
+                                      f'`{noise_type}`!')
+
+        # Set up bias.
+        if add_bias:
+            self.bias = nn.Parameter(torch.zeros(out_channels))
+            self.bscale = lr_mul
+        else:
+            self.bias = None
+
+        # Set up activation.
+        assert activation_type in ['linear', 'relu', 'lrelu']
+
+        # Set up style.
+        if use_style:
+            self.space_of_latent = 'W'
+            self.style = DenseLayer(in_channels=w_dim,
+                                    out_channels=out_channels * 2,
+                                    add_bias=True,
+                                    use_wscale=use_wscale,
+                                    wscale_gain=1.0,
+                                    lr_mul=1.0,
+                                    activation_type='linear')
+
+        if in_channels == 0:  # First layer.
+            self.const = nn.Parameter(
+                torch.ones(1, out_channels, resolution, resolution))
+            return
+
+        # Set up weight.
+        weight_shape = (out_channels, in_channels, kernel_size, kernel_size)
+        fan_in = kernel_size * kernel_size * in_channels
+        wscale = wscale_gain / np.sqrt(fan_in)
+        if use_wscale:
+            self.weight = nn.Parameter(torch.randn(*weight_shape) / lr_mul)
+            self.wscale = wscale * lr_mul
+        else:
+            self.weight = nn.Parameter(
+                torch.randn(*weight_shape) * wscale / lr_mul)
+            self.wscale = lr_mul
+
+        # Set up upsampling filter (if needed).
+        if scale_factor > 1:
+            assert filter_kernel is not None
+            kernel = np.array(filter_kernel, dtype=np.float32).reshape(1, -1)
+            kernel = kernel.T.dot(kernel)
+            kernel = kernel / np.sum(kernel)
+            kernel = kernel[np.newaxis, np.newaxis]
+            self.register_buffer('filter', torch.from_numpy(kernel))
+
+        if scale_factor > 1 and fused_scale:  # use transpose convolution.
+            self.stride = scale_factor
+        else:
+            self.stride = 1
+        self.padding = kernel_size // 2
+
+    def extra_repr(self):
+        return (f'in_ch={self.in_channels}, '
+                f'out_ch={self.out_channels}, '
+                f'ksize={self.kernel_size}, '
+                f'wscale_gain={self.wscale_gain:.3f}, '
+                f'bias={self.add_bias}, '
+                f'lr_mul={self.lr_mul:.3f}, '
+                f'upsample={self.scale_factor}, '
+                f'fused_scale={self.fused_scale}, '
+                f'upsample_filter={self.filter_kernel}, '
+                f'noise_type={self.noise_type}, '
+                f'act={self.activation_type}, '
+                f'use_style={self.use_style}')
+
+    def forward_style(self, w):
+        """Gets style code from the given input.
+
+        More specifically, if the input is from W-Space, it will be projected by
+        an affine transformation. If it is from the Style Space (Y-Space), no
+        operation is required.
+
+        NOTE: For codes from Y-Space, we use slicing to make sure the dimension
+        is correct, in case that the code is padded before fed into this layer.
+        """
+        space_of_latent = self.space_of_latent.upper()
+        if space_of_latent == 'W':
+            if w.ndim != 2 or w.shape[1] != self.w_dim:
+                raise ValueError(f'The input tensor should be with shape '
+                                 f'[batch_size, w_dim], where '
+                                 f'`w_dim` equals to {self.w_dim}!\n'
+                                 f'But `{w.shape}` is received!')
+            style = self.style(w)
+        elif space_of_latent == 'Y':
+            if w.ndim != 2 or w.shape[1] < self.out_channels * 2:
+                raise ValueError(f'The input tensor should be with shape '
+                                 f'[batch_size, y_dim], where '
+                                 f'`y_dim` equals to {self.out_channels * 2}!\n'
+                                 f'But `{w.shape}` is received!')
+            style = w[:, :self.out_channels * 2]
+        else:
+            raise NotImplementedError(f'Not implemented `space_of_latent`: '
+                                      f'`{space_of_latent}`!')
+        return style
+
+    def forward(self, x, w=None, noise_mode='const'):
+        if self.in_channels == 0:
+            assert x is None
+            x = self.const.repeat(w.shape[0], 1, 1, 1)
+        else:
+            weight = self.weight
+            if self.wscale != 1.0:
+                weight = weight * self.wscale
+
+            if self.scale_factor > 1 and self.fused_scale:
+                weight = F.pad(weight, (1, 1, 1, 1, 0, 0, 0, 0), 'constant', 0)
+                weight = (weight[:, :, 1:, 1:] + weight[:, :, :-1, 1:] +
+                          weight[:, :, 1:, :-1] + weight[:, :, :-1, :-1])
+                x = F.conv_transpose2d(x,
+                                       weight=weight.transpose(0, 1),
+                                       bias=None,
+                                       stride=self.stride,
+                                       padding=self.padding)
+            else:
+                if self.scale_factor > 1:
+                    up = self.scale_factor
+                    x = F.interpolate(x, scale_factor=up, mode='nearest')
+                x = F.conv2d(x,
+                             weight=weight,
+                             bias=None,
+                             stride=self.stride,
+                             padding=self.padding)
+
+            if self.scale_factor > 1:
+                # Disable `autocast` for customized autograd function.
+                # Please check reference:
+                # https://pytorch.org/docs/stable/notes/amp_examples.html#autocast-and-custom-autograd-functions
+                with autocast(enabled=False):
+                    f = self.filter.repeat(self.out_channels, 1, 1, 1)
+                    x = Blur.apply(x.float(), f)  # Always use FP32.
+
+        # Prepare noise.
+        noise_mode = noise_mode.lower()
+        if self.noise_type != 'none' and noise_mode != 'none':
+            if noise_mode == 'random':
+                noise = torch.randn(
+                    (x.shape[0], *self.noise.shape[1:]), device=x.device)
+            elif noise_mode == 'const':
+                noise = self.noise
+            else:
+                raise ValueError(f'Unknown noise mode `{noise_mode}`!')
+            x = x + noise * self.noise_strength
+
+        if self.bias is not None:
+            bias = self.bias
+            if self.bscale != 1.0:
+                bias = bias * self.bscale
+            x = x + bias.reshape(1, self.out_channels, 1, 1)
+
+        if self.activation_type == 'linear':
+            pass
+        elif self.activation_type == 'relu':
+            x = F.relu(x, inplace=True)
+        elif self.activation_type == 'lrelu':
+            x = F.leaky_relu(x, negative_slope=0.2, inplace=True)
+        else:
+            raise NotImplementedError(f'Not implemented activation type '
+                                      f'`{self.activation_type}`!')
+
+        if not self.use_style:
+            return x
+
+        # Instance normalization.
+        x = x - x.mean(dim=(2, 3), keepdim=True)
+        scale = (x.square().mean(dim=(2, 3), keepdim=True) + self.eps).rsqrt()
+        x = x * scale
+        # Style modulation.
+        style = self.forward_style(w)
+        style_split = style.unsqueeze(2).unsqueeze(3).chunk(2, dim=1)
+        x = x * (style_split[0] + 1) + style_split[1]
+
+        return x, style
+
+
+class DenseLayer(nn.Module):
+    """Implements the dense layer."""
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 add_bias,
+                 use_wscale,
+                 wscale_gain,
+                 lr_mul,
+                 activation_type):
+        """Initializes with layer settings.
+
+        Args:
+            in_channels: Number of channels of the input tensor.
+            out_channels: Number of channels of the output tensor.
+            add_bias: Whether to add bias onto the fully-connected result.
+            use_wscale: Whether to use weight scaling.
+            wscale_gain: Gain factor for weight scaling.
+            lr_mul: Learning multiplier for both weight and bias.
+            activation_type: Type of activation.
+        """
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.add_bias = add_bias
+        self.use_wscale = use_wscale
+        self.wscale_gain = wscale_gain
+        self.lr_mul = lr_mul
+        self.activation_type = activation_type
+
+        weight_shape = (out_channels, in_channels)
+        wscale = wscale_gain / np.sqrt(in_channels)
+        if use_wscale:
+            self.weight = nn.Parameter(torch.randn(*weight_shape) / lr_mul)
+            self.wscale = wscale * lr_mul
+        else:
+            self.weight = nn.Parameter(
+                torch.randn(*weight_shape) * wscale / lr_mul)
+            self.wscale = lr_mul
+
+        if add_bias:
+            self.bias = nn.Parameter(torch.zeros(out_channels))
+            self.bscale = lr_mul
+        else:
+            self.bias = None
+
+        assert activation_type in ['linear', 'relu', 'lrelu']
+
+    def extra_repr(self):
+        return (f'in_ch={self.in_channels}, '
+                f'out_ch={self.out_channels}, '
+                f'wscale_gain={self.wscale_gain:.3f}, '
+                f'bias={self.add_bias}, '
+                f'lr_mul={self.lr_mul:.3f}, '
+                f'act={self.activation_type}')
+
+    def forward(self, x):
+        if x.ndim != 2:
+            x = x.flatten(start_dim=1)
+
+        weight = self.weight
+        if self.wscale != 1.0:
+            weight = weight * self.wscale
+        bias = None
+        if self.bias is not None:
+            bias = self.bias
+            if self.bscale != 1.0:
+                bias = bias * self.bscale
+
+        x = F.linear(x, weight=weight, bias=bias)
+
+        if self.activation_type == 'linear':
+            pass
+        elif self.activation_type == 'relu':
+            x = F.relu(x, inplace=True)
+        elif self.activation_type == 'lrelu':
+            x = F.leaky_relu(x, negative_slope=0.2, inplace=True)
+        else:
+            raise NotImplementedError(f'Not implemented activation type '
+                                      f'`{self.activation_type}`!')
+
+        return x
+
+# pylint: enable=missing-function-docstring

models/test.py

@@ -0,0 +1,146 @@
+# python3.7
+"""Unit test for loading pre-trained models.
+
+Basically, this file tests whether the perceptual model (VGG16) and the
+inception model (InceptionV3), which are commonly used for loss computation and
+evaluation, have the expected behavior after loading pre-trained weights. In
+particular, we compare with the models from repo
+
+https://github.com/NVlabs/stylegan2-ada-pytorch
+"""
+
+import torch
+
+from models import build_model
+from utils.misc import download_url
+
+__all__ = ['test_model']
+
+_BATCH_SIZE = 4
+# pylint: disable=line-too-long
+_PERCEPTUAL_URL = 'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metrics/vgg16.pt'
+_INCEPTION_URL = 'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metrics/inception-2015-12-05.pt'
+# pylint: enable=line-too-long
+
+
+def test_model():
+    """Collects all model tests."""
+    torch.backends.cudnn.enabled = True
+    torch.backends.cudnn.allow_tf32 = False
+    torch.backends.cuda.matmul.allow_tf32 = False
+    torch.backends.cudnn.benchmark = False
+    torch.backends.cudnn.deterministic = True
+    print('========== Start Model Test ==========')
+    test_perceptual()
+    test_inception()
+    print('========== Finish Model Test ==========')
+
+
+def test_perceptual():
+    """Test the perceptual model."""
+    print('===== Testing Perceptual Model =====')
+
+    print('Build test model.')
+    model = build_model('PerceptualModel',
+                        use_torchvision=False,
+                        no_top=False,
+                        enable_lpips=True)
+
+    print('Build reference model.')
+    ref_model_path, _, = download_url(_PERCEPTUAL_URL)
+    with open(ref_model_path, 'rb') as f:
+        ref_model = torch.jit.load(f).eval().cuda()
+
+    print('Test performance: ')
+    for size in [224, 128, 256, 512, 1024]:
+        raw_img = torch.randint(0, 256, size=(_BATCH_SIZE, 3, size, size))
+        raw_img_comp = torch.randint(0, 256, size=(_BATCH_SIZE, 3, size, size))
+
+        # The test model requires input images to have range [-1, 1].
+        img = raw_img.to(torch.float32).cuda() / 127.5 - 1
+        img_comp = raw_img_comp.to(torch.float32).cuda() / 127.5 - 1
+        feat = model(img, resize_input=True, return_tensor='feature')
+        pred = model(img, resize_input=True, return_tensor='prediction')
+        lpips = model(img, img_comp, resize_input=False, return_tensor='lpips')
+        assert feat.shape == (_BATCH_SIZE, 4096)
+        assert pred.shape == (_BATCH_SIZE, 1000)
+        assert lpips.shape == (_BATCH_SIZE,)
+
+        # The reference model requires input images to have range [0, 255].
+        img = raw_img.to(torch.float32).cuda()
+        img_comp = raw_img_comp.to(torch.float32).cuda()
+        ref_feat = ref_model(img, resize_images=True, return_features=True)
+        ref_pred = ref_model(img, resize_images=True, return_features=False)
+        temp = ref_model(torch.cat([img, img_comp], dim=0),
+                         resize_images=False, return_lpips=True).chunk(2)
+        ref_lpips = (temp[0] - temp[1]).square().sum(dim=1, keepdim=False)
+        assert ref_feat.shape == (_BATCH_SIZE, 4096)
+        assert ref_pred.shape == (_BATCH_SIZE, 1000)
+        assert ref_lpips.shape == (_BATCH_SIZE,)
+
+        print(f'    Size {size}x{size}, feature (with resize):\n        '
+              f'mean: {(feat - ref_feat).abs().mean().item():.3e}, '
+              f'max: {(feat - ref_feat).abs().max().item():.3e}, '
+              f'ref_mean: {ref_feat.abs().mean().item():.3e}, '
+              f'ref_max: {ref_feat.abs().max().item():.3e}.')
+        print(f'    Size {size}x{size}, prediction (with resize):\n        '
+              f'mean: {(pred - ref_pred).abs().mean().item():.3e}, '
+              f'max: {(pred - ref_pred).abs().max().item():.3e}, '
+              f'ref_mean: {ref_pred.abs().mean().item():.3e}, '
+              f'ref_max: {ref_pred.abs().max().item():.3e}.')
+        print(f'    Size {size}x{size}, LPIPS (without resize):\n        '
+              f'mean: {(lpips - ref_lpips).abs().mean().item():.3e}, '
+              f'max: {(lpips - ref_lpips).abs().max().item():.3e}, '
+              f'ref_mean: {ref_lpips.abs().mean().item():.3e}, '
+              f'ref_max: {ref_lpips.abs().max().item():.3e}.')
+
+
+def test_inception():
+    """Test the inception model."""
+    print('===== Testing Inception Model =====')
+
+    print('Build test model.')
+    model = build_model('InceptionModel', align_tf=True)
+
+    print('Build reference model.')
+    ref_model_path, _, = download_url(_INCEPTION_URL)
+    with open(ref_model_path, 'rb') as f:
+        ref_model = torch.jit.load(f).eval().cuda()
+
+    print('Test performance: ')
+    for size in [299, 128, 256, 512, 1024]:
+        raw_img = torch.randint(0, 256, size=(_BATCH_SIZE, 3, size, size))
+
+        # The test model requires input images to have range [-1, 1].
+        img = raw_img.to(torch.float32).cuda() / 127.5 - 1
+        feat = model(img)
+        pred = model(img, output_predictions=True)
+        pred_nb = model(img, output_predictions=True, remove_logits_bias=True)
+        assert feat.shape == (_BATCH_SIZE, 2048)
+        assert pred.shape == (_BATCH_SIZE, 1008)
+        assert pred_nb.shape == (_BATCH_SIZE, 1008)
+
+        # The reference model requires input images to have range [0, 255].
+        img = raw_img.to(torch.float32).cuda()
+        ref_feat = ref_model(img, return_features=True)
+        ref_pred = ref_model(img)
+        ref_pred_nb = ref_model(img, no_output_bias=True)
+        assert ref_feat.shape == (_BATCH_SIZE, 2048)
+        assert ref_pred.shape == (_BATCH_SIZE, 1008)
+        assert ref_pred_nb.shape == (_BATCH_SIZE, 1008)
+
+        print(f'    Size {size}x{size}, feature:\n        '
+              f'mean: {(feat - ref_feat).abs().mean().item():.3e}, '
+              f'max: {(feat - ref_feat).abs().max().item():.3e}, '
+              f'ref_mean: {ref_feat.abs().mean().item():.3e}, '
+              f'ref_max: {ref_feat.abs().max().item():.3e}.')
+        print(f'    Size {size}x{size}, prediction:\n        '
+              f'mean: {(pred - ref_pred).abs().mean().item():.3e}, '
+              f'max: {(pred - ref_pred).abs().max().item():.3e}, '
+              f'ref_mean: {ref_pred.abs().mean().item():.3e}, '
+              f'ref_max: {ref_pred.abs().max().item():.3e}.')
+        print(f'    Size {size}x{size}, prediction (without bias):\n        '
+            f'mean: {(pred_nb - ref_pred_nb).abs().mean().item():.3e}, '
+            f'max: {(pred_nb - ref_pred_nb).abs().max().item():.3e}, '
+              f'ref_mean: {ref_pred_nb.abs().mean().item():.3e}, '
+              f'ref_max: {ref_pred_nb.abs().max().item():.3e}.')

models/utils/ops.py

@@ -0,0 +1,18 @@
+# python3.7
+"""Contains operators for neural networks."""
+
+import torch
+import torch.distributed as dist
+
+__all__ = ['all_gather']
+
+
+def all_gather(tensor):
+    """Gathers tensor from all devices and executes averaging."""
+    if not dist.is_initialized():
+        return tensor
+
+    world_size = dist.get_world_size()
+    tensor_list = [torch.ones_like(tensor) for _ in range(world_size)]
+    dist.all_gather(tensor_list, tensor, async_op=False)
+    return torch.stack(tensor_list, dim=0).mean(dim=0)

requirements/convert.txt

@@ -0,0 +1,11 @@
+torch==1.8.1
+tensorflow-gpu==1.15
+ninja==1.10.2
+scikit-video==1.1.11
+pillow==9.0.0
+opencv-python-headless==4.5.5.62
+requests
+bs4
+tqdm
+rich
+easydict

requirements/develop.txt

@@ -0,0 +1,3 @@
+bpytop  # Monitor system resources.
+gpustat  # Monitor GPU usage.
+pylint  # Check coding style.

requirements/minimal.txt

@@ -0,0 +1,21 @@
+torch==1.8.1+cu111 -f https://download.pytorch.org/whl/cu111/torch_stable.html
+torchvision==0.9.1+cu111 -f https://download.pytorch.org/whl/cu111/torch_stable.html
+tensorboard==2.7.0
+torch-tb-profiler==0.3.1
+ninja==1.10.2
+numpy==1.21.5
+scipy==1.7.3
+scikit-learn==1.0.2
+scikit-video==1.1.11
+pillow==9.0.0
+opencv-python-headless==4.5.5.62
+requests
+bs4
+tqdm
+rich
+click
+cloup
+psutil
+easydict
+lmdb
+matplotlib

synthesis.py

@@ -0,0 +1,178 @@
+# python3.7
+"""Script that synthesizes images with pre-trained models.
+
+Support StyleGAN2 and StyleGAN3.
+"""
+
+import os
+import argparse
+from tqdm import tqdm
+import numpy as np
+
+import torch
+from models import build_model
+from utils.visualizers.html_visualizer import HtmlVisualizer
+from utils.image_utils import save_image, resize_image
+from utils.image_utils import postprocess_image
+from utils.custom_utils import to_numpy
+
+
+def parse_args():
+    """Parses arguments."""
+    parser = argparse.ArgumentParser()
+    group = parser.add_argument_group('General options.')
+    group.add_argument('weight_path', type=str,
+                       help='Weight path to the pre-trained model.')
+    group.add_argument('--save_dir', type=str, default=None,
+                       help='Directory to save the results. If not specified, '
+                            'the results will be saved to '
+                            '`work_dirs/{TASK_SPECIFIC}/` by default.')
+    group.add_argument('--job', type=str, default='synthesize',
+                       help='Name for the job. (default: synthesize)')
+    group.add_argument('--seed', type=int, default=4,
+                       help='Seed for sampling. (default: 4)')
+    group.add_argument('--nums', type=int, default=100,
+                       help='Number of samples to synthesized. (default: 100)')
+    group.add_argument('--img_size', type=int, default=1024,
+                       help='Size of the synthesized images. (default: 1024)')
+    group.add_argument('--vis_size', type=int, default=256,
+                       help='Size of the visualize images. (default: 256)')
+    group.add_argument('--w_dim', type=int, default=512,
+                       help='Dimension of the latent w. (default: 512)')
+    group.add_argument('--batch_size', type=int, default=4,
+                       help='Batch size. (default: 4)')
+    group.add_argument('--save_jpg', action='store_true', default=False,
+                       help='Whether to save raw image. (default: False)')
+    group.add_argument('-d', '--data_name', type=str, default='ffhq',
+                       help='Name of the datasets. (default: ffhq)')
+    group.add_argument('--latent_path', type=str, default='',
+                       help='Path to the given latent codes. (default: None)')
+    group.add_argument('--trunc_psi', type=float, default=0.7,
+                       help='Psi factor used for truncation. (default: 0.7)')
+    group.add_argument('--trunc_layers', type=int, default=8,
+                       help='Number of layers to perform truncation.'
+                            ' (default: 8)')
+
+    group = parser.add_argument_group('StyleGAN2')
+    group.add_argument('--stylegan2', action='store_true',
+                       help='Whether or not using StyleGAN2. (default: False)')
+    group.add_argument('--scale_stylegan2', type=float, default=1.0,
+                       help='Scale for the number of channel fro stylegan2.')
+    group.add_argument('--randomize_noise', type=str, default='const',
+                       help='Noise type when synthesizing. (const or random)')
+
+    group = parser.add_argument_group('StyleGAN3')
+    group.add_argument('--stylegan3', action='store_true',
+                       help='Whether or not using StyleGAN3. (default: False)')
+    group.add_argument('--cfg', type=str, default='T',
+                       help='Config of the stylegan3 (T/R).')
+    group.add_argument('--scale_stylegan3r', type=float, default=2.0,
+                       help='Scale for the number of channel for stylegan3 R.')
+    group.add_argument('--scale_stylegan3t', type=float, default=1.0,
+                       help='Scale for the number of channel for stylegan3 T.')
+    group.add_argument('--tx', type=float, default=0,
+                       help='Translate X-coordinate. (default: 0.0)')
+    group.add_argument('--ty', type=float, default=0,
+                       help='Translate Y-coordinate. (default: 0.0)')
+    group.add_argument('--rotate', type=float, default=0,
+                       help='Rotation angle in degrees. (default: 0)')
+    return parser.parse_args()
+
+
+def main():
+    """Main function."""
+    args = parse_args()
+    # Parse model configuration.
+    assert (args.stylegan2 and not args.stylegan3) or \
+           (not args.stylegan2 and args.stylegan3)
+    job_disc = ''
+    if args.stylegan2:
+        config = dict(model_type='StyleGAN2Generator',
+                      resolution=args.img_size,
+                      w_dim=args.w_dim,
+                      fmaps_base=int(args.scale_stylegan2 * (32 << 10)),
+                      fmaps_max=512,)
+        job_disc += 'stylegan2'
+    else:
+        if args.stylegan3 and args.cfg == 'R':
+            config = dict(model_type='StyleGAN3Generator',
+                          resolution=args.img_size,
+                          w_dim=args.w_dim,
+                          fmaps_base=int(args.scale_stylegan3r * (32 << 10)),
+                          fmaps_max=1024,
+                          use_radial_filter=True,)
+            job_disc += 'stylegan3r'
+        elif args.stylegan3 and args.cfg == 'T':
+            config = dict(model_type='StyleGAN3Generator',
+                          resolution=args.img_size,
+                          w_dim=args.w_dim,
+                          fmaps_base=int(args.scale_stylegan3t * (32 << 10)),
+                          fmaps_max=512,
+                          use_radial_filter=False,
+                          kernel_size=3,)
+            job_disc += 'stylegan3t'
+        else:
+            raise TypeError(f'StyleGAN3 config type error, need `R/T`,'
+                            f' but got {args.cfg} instead.')
+
+    # Get work directory and job name.
+    save_dir = args.save_dir or f'work_dirs/{args.job}/{args.data_name}'
+    os.makedirs(save_dir, exist_ok=True)
+    job_name = f'seed_{args.seed}_num_{args.nums}_{job_disc}'
+    os.makedirs(f'{save_dir}/{job_name}', exist_ok=True)
+
+    # Build generation and get synthesis kwargs.
+    print('Building generator...')
+    generator = build_model(**config)
+    synthesis_kwargs = dict(trunc_psi=args.trunc_psi,
+                            trunc_layers=args.trunc_layers,)
+    # Load pre-trained weights.
+    checkpoint_path = args.weight_path
+    print(f'Loading checkpoint from `{checkpoint_path}` ...')
+    checkpoint = torch.load(checkpoint_path, map_location='cpu')['models']
+    if 'generator_smooth' in checkpoint:
+        generator.load_state_dict(checkpoint['generator_smooth'])
+    else:
+        generator.load_state_dict(checkpoint['generator'])
+    generator = generator.eval().cuda()
+    print('Finish loading checkpoint.')
+
+    np.random.seed(args.seed)
+    torch.manual_seed(args.seed)
+    if os.path.exists(args.latent_path):
+        latent_zs = np.load(args.latent_path)
+        latent_zs = latent_zs[:args.nums]
+    else:
+        latent_zs = np.random.randn(args.nums, generator.z_dim)
+    num_images = latent_zs.shape[0]
+    latent_zs = torch.from_numpy(latent_zs.astype(np.float32))
+    html = HtmlVisualizer(grid_size=num_images)
+    print(f'Synthesizing {num_images} images ...')
+    latent_ws = []
+    for batch_idx in tqdm(range(0, num_images, args.batch_size)):
+        latent_z = latent_zs[batch_idx:batch_idx + args.batch_size]
+        latent_z = latent_z.cuda()
+        with torch.no_grad():
+            g_outputs = generator(latent_z, **synthesis_kwargs)
+            g_image = to_numpy(g_outputs['image'])
+            images = postprocess_image(g_image)
+        for idx in range(images.shape[0]):
+            sub_idx = batch_idx + idx
+            img = images[idx]
+            row_idx, col_idx = divmod(sub_idx, html.num_cols)
+            image = resize_image(img, (args.vis_size, args.vis_size))
+            html.set_cell(row_idx, col_idx, image=image,
+                          text=f'Sample {sub_idx:06d}')
+            if args.save_jpg:
+                save_path = f'{save_dir}/{job_name}/{sub_idx:06d}.jpg'
+                save_image(save_path, img)
+        latent_ws.append(to_numpy(g_outputs['wp']))
+    latent_ws = np.concatenate(latent_ws, axis=0)
+    print(f'shape of the latent code: {latent_ws.shape}')
+    np.save(f'{save_dir}/{job_name}/latent_codes.npy', latent_ws)
+    html.save(f'{save_dir}/{job_name}.html')
+    print(f'Finish synthesizing {num_images} samples.')
+
+
+if __name__ == '__main__':
+    main()

third_party/stylegan2_official_ops/README.md

@@ -0,0 +1,28 @@
+# Operators for StyleGAN2
+
+All files in this directory are borrowed from repository [stylegan2-ada-pytorch](https://github.com/NVlabs/stylegan2-ada-pytorch). Basically, these files implement customized operators, which are faster than the native operators from PyTorch, especially for second-derivative computation, including
+
+- `bias_act.bias_act()`: Fuse adding bias and then performing activation as one operator.
+- `upfirdn2d.setup_filter()`: Set up the kernel used for filtering.
+- `upfirdn2d.filter2d()`: Filtering a 2D feature map with given kernel.
+- `upfirdn2d.upsample2d()`: Upsampling a 2D feature map.
+- `upfirdn2d.downsample2d()`: Downsampling a 2D feature map.
+- `upfirdn2d.upfirdn2d()`: Upsampling, filtering, and then downsampling a 2D feature map.
+- `conv2d_gradfix.conv2d()`: Convolutional layer, supporting arbitrarily high order gradients and fixing gradient when computing penalty.
+- `conv2d_gradfix.conv_transpose2d()`: Transposed convolutional layer, supporting arbitrarily high order gradients and fixing gradient when computing penalty.
+- `conv2d_resample.conv2d_resample()`: Wraps `upfirdn2d()` and `conv2d()` (or `conv_transpose2d()`). This is not used in our network implementation (*i.e.*, `models/stylegan2_generator.py` and `models/stylegan2_discriminator.py`)
+
+We make following slight modifications beyond disabling some lint warnings:
+
+- Line 25 of file `misc.py`: Use `EasyDict` from module `easydict` to replace that from `dnnlib` from [stylegan2-ada-pytorch](https://github.com/NVlabs/stylegan2-ada-pytorch).
+- Line 35 of file `custom_ops.py`: Disable log message when setting up customized operators.
+- Line 53/89 of file `custom_ops.py`: Add necessary CUDA compiler path. (***NOTE**: If your cuda binary does not locate at `/usr/local/cuda/bin`, please specify in function `_find_compiler_bindir_posix()`.*)
+- Line 24 of file `bias_act.py`: Use `EasyDict` from module `easydict` to replace that from `dnnlib` from [stylegan2-ada-pytorch](https://github.com/NVlabs/stylegan2-ada-pytorch).
+- Line 32 of file `grid_sample_gradfix.py`: Enable customized grid sampling operator by default.
+- Line 36 of file `grid_sample_gradfix.py`:  Use `impl` to disable customized grid sample operator.
+- Line 33 of file `conv2d_gradfix.py`: Enable customized convolution operators by default.
+- Line 46/51 of file `conv2d_gradfix.py`: Use `impl` to disable customized convolution operators.
+- Line 36/66 of file `conv2d_resample.py`: Use `impl` to disable customized convolution operators.
+- Line 23 of file `fma.py`: Use `impl` to disable customized add-multiply operator.
+
+Please use `ref` or `cuda` to choose which implementation to use. `ref` refers to native PyTorch operators while `cuda` refers to the customized operators from the official repository. `cuda` is used by default.

third_party/stylegan2_official_ops/bias_act.cpp

@@ -0,0 +1,99 @@
+// 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
+// and any modifications thereto.  Any use, reproduction, disclosure or
+// distribution of this software and related documentation without an express
+// license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+#include <torch/extension.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include "bias_act.h"
+
+//------------------------------------------------------------------------
+
+static bool has_same_layout(torch::Tensor x, torch::Tensor y)
+{
+    if (x.dim() != y.dim())
+        return false;
+    for (int64_t i = 0; i < x.dim(); i++)
+    {
+        if (x.size(i) != y.size(i))
+            return false;
+        if (x.size(i) >= 2 && x.stride(i) != y.stride(i))
+            return false;
+    }
+    return true;
+}
+
+//------------------------------------------------------------------------
+
+static torch::Tensor bias_act(torch::Tensor x, torch::Tensor b, torch::Tensor xref, torch::Tensor yref, torch::Tensor dy, int grad, int dim, int act, float alpha, float gain, float clamp)
+{
+    // Validate arguments.
+    TORCH_CHECK(x.is_cuda(), "x must reside on CUDA device");
+    TORCH_CHECK(b.numel() == 0 || (b.dtype() == x.dtype() && b.device() == x.device()), "b must have the same dtype and device as x");
+    TORCH_CHECK(xref.numel() == 0 || (xref.sizes() == x.sizes() && xref.dtype() == x.dtype() && xref.device() == x.device()), "xref must have the same shape, dtype, and device as x");
+    TORCH_CHECK(yref.numel() == 0 || (yref.sizes() == x.sizes() && yref.dtype() == x.dtype() && yref.device() == x.device()), "yref must have the same shape, dtype, and device as x");
+    TORCH_CHECK(dy.numel() == 0 || (dy.sizes() == x.sizes() && dy.dtype() == x.dtype() && dy.device() == x.device()), "dy must have the same dtype and device as x");
+    TORCH_CHECK(x.numel() <= INT_MAX, "x is too large");
+    TORCH_CHECK(b.dim() == 1, "b must have rank 1");
+    TORCH_CHECK(b.numel() == 0 || (dim >= 0 && dim < x.dim()), "dim is out of bounds");
+    TORCH_CHECK(b.numel() == 0 || b.numel() == x.size(dim), "b has wrong number of elements");
+    TORCH_CHECK(grad >= 0, "grad must be non-negative");
+
+    // Validate layout.
+    TORCH_CHECK(x.is_non_overlapping_and_dense(), "x must be non-overlapping and dense");
+    TORCH_CHECK(b.is_contiguous(), "b must be contiguous");
+    TORCH_CHECK(xref.numel() == 0 || has_same_layout(xref, x), "xref must have the same layout as x");
+    TORCH_CHECK(yref.numel() == 0 || has_same_layout(yref, x), "yref must have the same layout as x");
+    TORCH_CHECK(dy.numel() == 0 || has_same_layout(dy, x), "dy must have the same layout as x");
+
+    // Create output tensor.
+    const at::cuda::OptionalCUDAGuard device_guard(device_of(x));
+    torch::Tensor y = torch::empty_like(x);
+    TORCH_CHECK(has_same_layout(y, x), "y must have the same layout as x");
+
+    // Initialize CUDA kernel parameters.
+    bias_act_kernel_params p;
+    p.x     = x.data_ptr();
+    p.b     = (b.numel()) ? b.data_ptr() : NULL;
+    p.xref  = (xref.numel()) ? xref.data_ptr() : NULL;
+    p.yref  = (yref.numel()) ? yref.data_ptr() : NULL;
+    p.dy    = (dy.numel()) ? dy.data_ptr() : NULL;
+    p.y     = y.data_ptr();
+    p.grad  = grad;
+    p.act   = act;
+    p.alpha = alpha;
+    p.gain  = gain;
+    p.clamp = clamp;
+    p.sizeX = (int)x.numel();
+    p.sizeB = (int)b.numel();
+    p.stepB = (b.numel()) ? (int)x.stride(dim) : 1;
+
+    // Choose CUDA kernel.
+    void* kernel;
+    AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), "upfirdn2d_cuda", [&]
+    {
+        kernel = choose_bias_act_kernel<scalar_t>(p);
+    });
+    TORCH_CHECK(kernel, "no CUDA kernel found for the specified activation func");
+
+    // Launch CUDA kernel.
+    p.loopX = 4;
+    int blockSize = 4 * 32;
+    int gridSize = (p.sizeX - 1) / (p.loopX * blockSize) + 1;
+    void* args[] = {&p};
+    AT_CUDA_CHECK(cudaLaunchKernel(kernel, gridSize, blockSize, args, 0, at::cuda::getCurrentCUDAStream()));
+    return y;
+}
+
+//------------------------------------------------------------------------
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
+{
+    m.def("bias_act", &bias_act);
+}
+
+//------------------------------------------------------------------------

third_party/stylegan2_official_ops/bias_act.cu

@@ -0,0 +1,173 @@
+// 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
+// and any modifications thereto.  Any use, reproduction, disclosure or
+// distribution of this software and related documentation without an express
+// license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+#include <c10/util/Half.h>
+#include "bias_act.h"
+
+//------------------------------------------------------------------------
+// Helpers.
+
+template <class T> struct InternalType;
+template <> struct InternalType<double>     { typedef double scalar_t; };
+template <> struct InternalType<float>      { typedef float  scalar_t; };
+template <> struct InternalType<c10::Half>  { typedef float  scalar_t; };
+
+//------------------------------------------------------------------------
+// CUDA kernel.
+
+template <class T, int A>
+__global__ void bias_act_kernel(bias_act_kernel_params p)
+{
+    typedef typename InternalType<T>::scalar_t scalar_t;
+    int G                 = p.grad;
+    scalar_t alpha        = (scalar_t)p.alpha;
+    scalar_t gain         = (scalar_t)p.gain;
+    scalar_t clamp        = (scalar_t)p.clamp;
+    scalar_t one          = (scalar_t)1;
+    scalar_t two          = (scalar_t)2;
+    scalar_t expRange     = (scalar_t)80;
+    scalar_t halfExpRange = (scalar_t)40;
+    scalar_t seluScale    = (scalar_t)1.0507009873554804934193349852946;
+    scalar_t seluAlpha    = (scalar_t)1.6732632423543772848170429916717;
+
+    // Loop over elements.
+    int xi = blockIdx.x * p.loopX * blockDim.x + threadIdx.x;
+    for (int loopIdx = 0; loopIdx < p.loopX && xi < p.sizeX; loopIdx++, xi += blockDim.x)
+    {
+        // Load.
+        scalar_t x = (scalar_t)((const T*)p.x)[xi];
+        scalar_t b = (p.b) ? (scalar_t)((const T*)p.b)[(xi / p.stepB) % p.sizeB] : 0;
+        scalar_t xref = (p.xref) ? (scalar_t)((const T*)p.xref)[xi] : 0;
+        scalar_t yref = (p.yref) ? (scalar_t)((const T*)p.yref)[xi] : 0;
+        scalar_t dy = (p.dy) ? (scalar_t)((const T*)p.dy)[xi] : one;
+        scalar_t yy = (gain != 0) ? yref / gain : 0;
+        scalar_t y = 0;
+
+        // Apply bias.
+        ((G == 0) ? x : xref) += b;
+
+        // linear
+        if (A == 1)
+        {
+            if (G == 0) y = x;
+            if (G == 1) y = x;
+        }
+
+        // relu
+        if (A == 2)
+        {
+            if (G == 0) y = (x > 0) ? x : 0;
+            if (G == 1) y = (yy > 0) ? x : 0;
+        }
+
+        // lrelu
+        if (A == 3)
+        {
+            if (G == 0) y = (x > 0) ? x : x * alpha;
+            if (G == 1) y = (yy > 0) ? x : x * alpha;
+        }
+
+        // tanh
+        if (A == 4)
+        {
+            if (G == 0) { scalar_t c = exp(x); scalar_t d = one / c; y = (x < -expRange) ? -one : (x > expRange) ? one : (c - d) / (c + d); }
+            if (G == 1) y = x * (one - yy * yy);
+            if (G == 2) y = x * (one - yy * yy) * (-two * yy);
+        }
+
+        // sigmoid
+        if (A == 5)
+        {
+            if (G == 0) y = (x < -expRange) ? 0 : one / (exp(-x) + one);
+            if (G == 1) y = x * yy * (one - yy);
+            if (G == 2) y = x * yy * (one - yy) * (one - two * yy);
+        }
+
+        // elu
+        if (A == 6)
+        {
+            if (G == 0) y = (x >= 0) ? x : exp(x) - one;
+            if (G == 1) y = (yy >= 0) ? x : x * (yy + one);
+            if (G == 2) y = (yy >= 0) ? 0 : x * (yy + one);
+        }
+
+        // selu
+        if (A == 7)
+        {
+            if (G == 0) y = (x >= 0) ? seluScale * x : (seluScale * seluAlpha) * (exp(x) - one);
+            if (G == 1) y = (yy >= 0) ? x * seluScale : x * (yy + seluScale * seluAlpha);
+            if (G == 2) y = (yy >= 0) ? 0 : x * (yy + seluScale * seluAlpha);
+        }
+
+        // softplus
+        if (A == 8)
+        {
+            if (G == 0) y = (x > expRange) ? x : log(exp(x) + one);
+            if (G == 1) y = x * (one - exp(-yy));
+            if (G == 2) { scalar_t c = exp(-yy); y = x * c * (one - c); }
+        }
+
+        // swish
+        if (A == 9)
+        {
+            if (G == 0)
+                y = (x < -expRange) ? 0 : x / (exp(-x) + one);
+            else
+            {
+                scalar_t c = exp(xref);
+                scalar_t d = c + one;
+                if (G == 1)
+                    y = (xref > halfExpRange) ? x : x * c * (xref + d) / (d * d);
+                else
+                    y = (xref > halfExpRange) ? 0 : x * c * (xref * (two - d) + two * d) / (d * d * d);
+                yref = (xref < -expRange) ? 0 : xref / (exp(-xref) + one) * gain;
+            }
+        }
+
+        // Apply gain.
+        y *= gain * dy;
+
+        // Clamp.
+        if (clamp >= 0)
+        {
+            if (G == 0)
+                y = (y > -clamp & y < clamp) ? y : (y >= 0) ? clamp : -clamp;
+            else
+                y = (yref > -clamp & yref < clamp) ? y : 0;
+        }
+
+        // Store.
+        ((T*)p.y)[xi] = (T)y;
+    }
+}
+
+//------------------------------------------------------------------------
+// CUDA kernel selection.
+
+template <class T> void* choose_bias_act_kernel(const bias_act_kernel_params& p)
+{
+    if (p.act == 1) return (void*)bias_act_kernel<T, 1>;
+    if (p.act == 2) return (void*)bias_act_kernel<T, 2>;
+    if (p.act == 3) return (void*)bias_act_kernel<T, 3>;
+    if (p.act == 4) return (void*)bias_act_kernel<T, 4>;
+    if (p.act == 5) return (void*)bias_act_kernel<T, 5>;
+    if (p.act == 6) return (void*)bias_act_kernel<T, 6>;
+    if (p.act == 7) return (void*)bias_act_kernel<T, 7>;
+    if (p.act == 8) return (void*)bias_act_kernel<T, 8>;
+    if (p.act == 9) return (void*)bias_act_kernel<T, 9>;
+    return NULL;
+}
+
+//------------------------------------------------------------------------
+// Template specializations.
+
+template void* choose_bias_act_kernel<double>       (const bias_act_kernel_params& p);
+template void* choose_bias_act_kernel<float>        (const bias_act_kernel_params& p);
+template void* choose_bias_act_kernel<c10::Half>    (const bias_act_kernel_params& p);
+
+//------------------------------------------------------------------------

third_party/stylegan2_official_ops/bias_act.h

@@ -0,0 +1,38 @@
+// 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
+// and any modifications thereto.  Any use, reproduction, disclosure or
+// distribution of this software and related documentation without an express
+// license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+//------------------------------------------------------------------------
+// CUDA kernel parameters.
+
+struct bias_act_kernel_params
+{
+    const void* x;      // [sizeX]
+    const void* b;      // [sizeB] or NULL
+    const void* xref;   // [sizeX] or NULL
+    const void* yref;   // [sizeX] or NULL
+    const void* dy;     // [sizeX] or NULL
+    void*       y;      // [sizeX]
+
+    int         grad;
+    int         act;
+    float       alpha;
+    float       gain;
+    float       clamp;
+
+    int         sizeX;
+    int         sizeB;
+    int         stepB;
+    int         loopX;
+};
+
+//------------------------------------------------------------------------
+// CUDA kernel selection.
+
+template <class T> void* choose_bias_act_kernel(const bias_act_kernel_params& p);
+
+//------------------------------------------------------------------------

third_party/stylegan2_official_ops/bias_act.py

@@ -0,0 +1,227 @@
+# python3.7
+
+# 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
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+"""Custom ops to fuse bias and activation as one operator, which is efficient.
+
+Please refer to https://github.com/NVlabs/stylegan2-ada-pytorch
+"""
+
+# pylint: disable=line-too-long
+# pylint: disable=missing-class-docstring
+# pylint: disable=global-statement
+# pylint: disable=bare-except
+
+import os
+import warnings
+import traceback
+from easydict import EasyDict
+import numpy as np
+import torch
+
+from . import custom_ops
+from . import misc
+
+#----------------------------------------------------------------------------
+
+activation_funcs = {
+    'linear':   EasyDict(func=lambda x, **_:         x,                                          def_alpha=0,    def_gain=1,             cuda_idx=1, ref='',  has_2nd_grad=False),
+    'relu':     EasyDict(func=lambda x, **_:         torch.nn.functional.relu(x),                def_alpha=0,    def_gain=np.sqrt(2),    cuda_idx=2, ref='y', has_2nd_grad=False),
+    'lrelu':    EasyDict(func=lambda x, alpha, **_:  torch.nn.functional.leaky_relu(x, alpha),   def_alpha=0.2,  def_gain=np.sqrt(2),    cuda_idx=3, ref='y', has_2nd_grad=False),
+    'tanh':     EasyDict(func=lambda x, **_:         torch.tanh(x),                              def_alpha=0,    def_gain=1,             cuda_idx=4, ref='y', has_2nd_grad=True),
+    'sigmoid':  EasyDict(func=lambda x, **_:         torch.sigmoid(x),                           def_alpha=0,    def_gain=1,             cuda_idx=5, ref='y', has_2nd_grad=True),
+    'elu':      EasyDict(func=lambda x, **_:         torch.nn.functional.elu(x),                 def_alpha=0,    def_gain=1,             cuda_idx=6, ref='y', has_2nd_grad=True),
+    'selu':     EasyDict(func=lambda x, **_:         torch.nn.functional.selu(x),                def_alpha=0,    def_gain=1,             cuda_idx=7, ref='y', has_2nd_grad=True),
+    'softplus': EasyDict(func=lambda x, **_:         torch.nn.functional.softplus(x),            def_alpha=0,    def_gain=1,             cuda_idx=8, ref='y', has_2nd_grad=True),
+    'swish':    EasyDict(func=lambda x, **_:         torch.sigmoid(x) * x,                       def_alpha=0,    def_gain=np.sqrt(2),    cuda_idx=9, ref='x', has_2nd_grad=True),
+}
+
+#----------------------------------------------------------------------------
+
+_inited = False
+_plugin = None
+_null_tensor = torch.empty([0])
+
+def _init():
+    global _inited, _plugin
+    if not _inited:
+        _inited = True
+        sources = ['bias_act.cpp', 'bias_act.cu']
+        sources = [os.path.join(os.path.dirname(__file__), s) for s in sources]
+        try:
+            _plugin = custom_ops.get_plugin('bias_act_plugin', sources=sources, extra_cuda_cflags=['--use_fast_math'])
+        except:
+            warnings.warn('Failed to build CUDA kernels for bias_act. Falling back to slow reference implementation. Details:\n\n' + traceback.format_exc())
+    return _plugin is not None
+
+#----------------------------------------------------------------------------
+
+def bias_act(x, b=None, dim=1, act='linear', alpha=None, gain=None, clamp=None, impl='cuda'):
+    r"""Fused bias and activation function.
+
+    Adds bias `b` to activation tensor `x`, evaluates activation function `act`,
+    and scales the result by `gain`. Each of the steps is optional. In most cases,
+    the fused op is considerably more efficient than performing the same calculation
+    using standard PyTorch ops. It supports first and second order gradients,
+    but not third order gradients.
+
+    Args:
+        x:      Input activation tensor. Can be of any shape.
+        b:      Bias vector, or `None` to disable. Must be a 1D tensor of the same type
+                as `x`. The shape must be known, and it must match the dimension of `x`
+                corresponding to `dim`.
+        dim:    The dimension in `x` corresponding to the elements of `b`.
+                The value of `dim` is ignored if `b` is not specified.
+        act:    Name of the activation function to evaluate, or `"linear"` to disable.
+                Can be e.g. `"relu"`, `"lrelu"`, `"tanh"`, `"sigmoid"`, `"swish"`, etc.
+                See `activation_funcs` for a full list. `None` is not allowed.
+        alpha:  Shape parameter for the activation function, or `None` to use the default.
+        gain:   Scaling factor for the output tensor, or `None` to use default.
+                See `activation_funcs` for the default scaling of each activation function.
+                If unsure, consider specifying 1.
+        clamp:  Clamp the output values to `[-clamp, +clamp]`, or `None` to disable
+                the clamping (default).
+        impl:   Name of the implementation to use. Can be `"ref"` or `"cuda"` (default).
+
+    Returns:
+        Tensor of the same shape and datatype as `x`.
+    """
+    assert isinstance(x, torch.Tensor)
+    assert impl in ['ref', 'cuda']
+    if impl == 'cuda' and x.device.type == 'cuda' and _init():
+        return _bias_act_cuda(dim=dim, act=act, alpha=alpha, gain=gain, clamp=clamp).apply(x, b)
+    return _bias_act_ref(x=x, b=b, dim=dim, act=act, alpha=alpha, gain=gain, clamp=clamp)
+
+#----------------------------------------------------------------------------
+
+@misc.profiled_function
+def _bias_act_ref(x, b=None, dim=1, act='linear', alpha=None, gain=None, clamp=None):
+    """Slow reference implementation of `bias_act()` using standard TensorFlow ops.
+    """
+    assert isinstance(x, torch.Tensor)
+    assert clamp is None or clamp >= 0
+    spec = activation_funcs[act]
+    alpha = float(alpha if alpha is not None else spec.def_alpha)
+    gain = float(gain if gain is not None else spec.def_gain)
+    clamp = float(clamp if clamp is not None else -1)
+
+    # Add bias.
+    if b is not None:
+        assert isinstance(b, torch.Tensor) and b.ndim == 1
+        assert 0 <= dim < x.ndim
+        assert b.shape[0] == x.shape[dim]
+        x = x + b.reshape([-1 if i == dim else 1 for i in range(x.ndim)])
+
+    # Evaluate activation function.
+    alpha = float(alpha)
+    x = spec.func(x, alpha=alpha)
+
+    # Scale by gain.
+    gain = float(gain)
+    if gain != 1:
+        x = x * gain
+
+    # Clamp.
+    if clamp >= 0:
+        x = x.clamp(-clamp, clamp)
+    return x
+
+#----------------------------------------------------------------------------
+
+_bias_act_cuda_cache = dict()
+
+def _bias_act_cuda(dim=1, act='linear', alpha=None, gain=None, clamp=None):
+    """Fast CUDA implementation of `bias_act()` using custom ops.
+    """
+    # Parse arguments.
+    assert clamp is None or clamp >= 0
+    spec = activation_funcs[act]
+    alpha = float(alpha if alpha is not None else spec.def_alpha)
+    gain = float(gain if gain is not None else spec.def_gain)
+    clamp = float(clamp if clamp is not None else -1)
+
+    # Lookup from cache.
+    key = (dim, act, alpha, gain, clamp)
+    if key in _bias_act_cuda_cache:
+        return _bias_act_cuda_cache[key]
+
+    # Forward op.
+    class BiasActCuda(torch.autograd.Function):
+        @staticmethod
+        def forward(ctx, x, b): # pylint: disable=arguments-differ
+            ctx.memory_format = torch.channels_last if x.ndim > 2 and x.stride()[1] == 1 else torch.contiguous_format
+            x = x.contiguous(memory_format=ctx.memory_format)
+            b = b.contiguous() if b is not None else _null_tensor
+            y = x
+            if act != 'linear' or gain != 1 or clamp >= 0 or b is not _null_tensor:
+                y = _plugin.bias_act(x, b, _null_tensor, _null_tensor, _null_tensor, 0, dim, spec.cuda_idx, alpha, gain, clamp)
+            ctx.save_for_backward(
+                x if 'x' in spec.ref or spec.has_2nd_grad else _null_tensor,
+                b if 'x' in spec.ref or spec.has_2nd_grad else _null_tensor,
+                y if 'y' in spec.ref else _null_tensor)
+            return y
+
+        @staticmethod
+        def backward(ctx, dy): # pylint: disable=arguments-differ
+            dy = dy.contiguous(memory_format=ctx.memory_format)
+            x, b, y = ctx.saved_tensors
+            dx = None
+            db = None
+
+            if ctx.needs_input_grad[0] or ctx.needs_input_grad[1]:
+                dx = dy
+                if act != 'linear' or gain != 1 or clamp >= 0:
+                    dx = BiasActCudaGrad.apply(dy, x, b, y)
+
+            if ctx.needs_input_grad[1]:
+                db = dx.sum([i for i in range(dx.ndim) if i != dim])
+
+            return dx, db
+
+    # Backward op.
+    class BiasActCudaGrad(torch.autograd.Function):
+        @staticmethod
+        def forward(ctx, dy, x, b, y): # pylint: disable=arguments-differ
+            ctx.memory_format = torch.channels_last if dy.ndim > 2 and dy.stride()[1] == 1 else torch.contiguous_format
+            dx = _plugin.bias_act(dy, b, x, y, _null_tensor, 1, dim, spec.cuda_idx, alpha, gain, clamp)
+            ctx.save_for_backward(
+                dy if spec.has_2nd_grad else _null_tensor,
+                x, b, y)
+            return dx
+
+        @staticmethod
+        def backward(ctx, d_dx): # pylint: disable=arguments-differ
+            d_dx = d_dx.contiguous(memory_format=ctx.memory_format)
+            dy, x, b, y = ctx.saved_tensors
+            d_dy = None
+            d_x = None
+            d_b = None
+            d_y = None
+
+            if ctx.needs_input_grad[0]:
+                d_dy = BiasActCudaGrad.apply(d_dx, x, b, y)
+
+            if spec.has_2nd_grad and (ctx.needs_input_grad[1] or ctx.needs_input_grad[2]):
+                d_x = _plugin.bias_act(d_dx, b, x, y, dy, 2, dim, spec.cuda_idx, alpha, gain, clamp)
+
+            if spec.has_2nd_grad and ctx.needs_input_grad[2]:
+                d_b = d_x.sum([i for i in range(d_x.ndim) if i != dim])
+
+            return d_dy, d_x, d_b, d_y
+
+    # Add to cache.
+    _bias_act_cuda_cache[key] = BiasActCuda
+    return BiasActCuda
+
+#----------------------------------------------------------------------------
+
+# pylint: enable=line-too-long
+# pylint: enable=missing-class-docstring
+# pylint: enable=global-statement
+# pylint: enable=bare-except

third_party/stylegan2_official_ops/conv2d_gradfix.py

@@ -0,0 +1,189 @@
+# python3.7
+
+# 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
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+"""Custom replacement for convolution operators.
+
+Operators in this file support arbitrarily high order gradients with zero
+performance penalty. Please set `impl` as `cuda` to use faster customized
+operators, OR as `ref` to use native `torch.nn.functional.conv2d` and
+`torch.nn.functional.conv_transpose2d`.
+
+Please refer to https://github.com/NVlabs/stylegan2-ada-pytorch
+"""
+
+# pylint: disable=redefined-builtin
+# pylint: disable=arguments-differ
+# pylint: disable=protected-access
+# pylint: disable=line-too-long
+# pylint: disable=global-statement
+# pylint: disable=missing-class-docstring
+# pylint: disable=missing-function-docstring
+
+import warnings
+import contextlib
+import torch
+
+enabled = True                     # Enable the custom op by setting this to true.
+weight_gradients_disabled = False  # Forcefully disable computation of gradients with respect to the weights.
+
+@contextlib.contextmanager
+def no_weight_gradients():
+    global weight_gradients_disabled
+    old = weight_gradients_disabled
+    weight_gradients_disabled = True
+    yield
+    weight_gradients_disabled = old
+
+#----------------------------------------------------------------------------
+
+def conv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1, impl='cuda'):
+    if impl == 'cuda' and _should_use_custom_op(input):
+        return _conv2d_gradfix(transpose=False, weight_shape=weight.shape, stride=stride, padding=padding, output_padding=0, dilation=dilation, groups=groups).apply(input, weight, bias)
+    return torch.nn.functional.conv2d(input=input, weight=weight, bias=bias, stride=stride, padding=padding, dilation=dilation, groups=groups)
+
+def conv_transpose2d(input, weight, bias=None, stride=1, padding=0, output_padding=0, groups=1, dilation=1, impl='cuda'):
+    if impl == 'cuda' and _should_use_custom_op(input):
+        return _conv2d_gradfix(transpose=True, weight_shape=weight.shape, stride=stride, padding=padding, output_padding=output_padding, groups=groups, dilation=dilation).apply(input, weight, bias)
+    return torch.nn.functional.conv_transpose2d(input=input, weight=weight, bias=bias, stride=stride, padding=padding, output_padding=output_padding, groups=groups, dilation=dilation)
+
+#----------------------------------------------------------------------------
+
+def _should_use_custom_op(input):
+    assert isinstance(input, torch.Tensor)
+    if (not enabled) or (not torch.backends.cudnn.enabled):
+        return False
+    if input.device.type != 'cuda':
+        return False
+    if any(torch.__version__.startswith(x) for x in ['1.7.', '1.8.', '1.9']):
+        return True
+    warnings.warn(f'conv2d_gradfix not supported on PyTorch {torch.__version__}. Falling back to torch.nn.functional.conv2d().')
+    return False
+
+def _tuple_of_ints(xs, ndim):
+    xs = tuple(xs) if isinstance(xs, (tuple, list)) else (xs,) * ndim
+    assert len(xs) == ndim
+    assert all(isinstance(x, int) for x in xs)
+    return xs
+
+#----------------------------------------------------------------------------
+
+_conv2d_gradfix_cache = dict()
+
+def _conv2d_gradfix(transpose, weight_shape, stride, padding, output_padding, dilation, groups):
+    # Parse arguments.
+    ndim = 2
+    weight_shape = tuple(weight_shape)
+    stride = _tuple_of_ints(stride, ndim)
+    padding = _tuple_of_ints(padding, ndim)
+    output_padding = _tuple_of_ints(output_padding, ndim)
+    dilation = _tuple_of_ints(dilation, ndim)
+
+    # Lookup from cache.
+    key = (transpose, weight_shape, stride, padding, output_padding, dilation, groups)
+    if key in _conv2d_gradfix_cache:
+        return _conv2d_gradfix_cache[key]
+
+    # Validate arguments.
+    assert groups >= 1
+    assert len(weight_shape) == ndim + 2
+    assert all(stride[i] >= 1 for i in range(ndim))
+    assert all(padding[i] >= 0 for i in range(ndim))
+    assert all(dilation[i] >= 0 for i in range(ndim))
+    if not transpose:
+        assert all(output_padding[i] == 0 for i in range(ndim))
+    else: # transpose
+        assert all(0 <= output_padding[i] < max(stride[i], dilation[i]) for i in range(ndim))
+
+    # Helpers.
+    common_kwargs = dict(stride=stride, padding=padding, dilation=dilation, groups=groups)
+    def calc_output_padding(input_shape, output_shape):
+        if transpose:
+            return [0, 0]
+        return [
+            input_shape[i + 2]
+            - (output_shape[i + 2] - 1) * stride[i]
+            - (1 - 2 * padding[i])
+            - dilation[i] * (weight_shape[i + 2] - 1)
+            for i in range(ndim)
+        ]
+
+    # Forward & backward.
+    class Conv2d(torch.autograd.Function):
+        @staticmethod
+        def forward(ctx, input, weight, bias):
+            assert weight.shape == weight_shape
+            if not transpose:
+                output = torch.nn.functional.conv2d(input=input, weight=weight, bias=bias, **common_kwargs)
+            else: # transpose
+                output = torch.nn.functional.conv_transpose2d(input=input, weight=weight, bias=bias, output_padding=output_padding, **common_kwargs)
+            ctx.save_for_backward(input, weight)
+            return output
+
+        @staticmethod
+        def backward(ctx, grad_output):
+            input, weight = ctx.saved_tensors
+            grad_input = None
+            grad_weight = None
+            grad_bias = None
+
+            if ctx.needs_input_grad[0]:
+                p = calc_output_padding(input_shape=input.shape, output_shape=grad_output.shape)
+                grad_input = _conv2d_gradfix(transpose=(not transpose), weight_shape=weight_shape, output_padding=p, **common_kwargs).apply(grad_output, weight, None)
+                assert grad_input.shape == input.shape
+
+            if ctx.needs_input_grad[1] and not weight_gradients_disabled:
+                grad_weight = Conv2dGradWeight.apply(grad_output, input)
+                assert grad_weight.shape == weight_shape
+
+            if ctx.needs_input_grad[2]:
+                grad_bias = grad_output.sum([0, 2, 3])
+
+            return grad_input, grad_weight, grad_bias
+
+    # Gradient with respect to the weights.
+    class Conv2dGradWeight(torch.autograd.Function):
+        @staticmethod
+        def forward(ctx, grad_output, input):
+            op = torch._C._jit_get_operation('aten::cudnn_convolution_backward_weight' if not transpose else 'aten::cudnn_convolution_transpose_backward_weight')
+            flags = [torch.backends.cudnn.benchmark, torch.backends.cudnn.deterministic, torch.backends.cudnn.allow_tf32]
+            grad_weight = op(weight_shape, grad_output, input, padding, stride, dilation, groups, *flags)
+            assert grad_weight.shape == weight_shape
+            ctx.save_for_backward(grad_output, input)
+            return grad_weight
+
+        @staticmethod
+        def backward(ctx, grad2_grad_weight):
+            grad_output, input = ctx.saved_tensors
+            grad2_grad_output = None
+            grad2_input = None
+
+            if ctx.needs_input_grad[0]:
+                grad2_grad_output = Conv2d.apply(input, grad2_grad_weight, None)
+                assert grad2_grad_output.shape == grad_output.shape
+
+            if ctx.needs_input_grad[1]:
+                p = calc_output_padding(input_shape=input.shape, output_shape=grad_output.shape)
+                grad2_input = _conv2d_gradfix(transpose=(not transpose), weight_shape=weight_shape, output_padding=p, **common_kwargs).apply(grad_output, grad2_grad_weight, None)
+                assert grad2_input.shape == input.shape
+
+            return grad2_grad_output, grad2_input
+
+    _conv2d_gradfix_cache[key] = Conv2d
+    return Conv2d
+
+#----------------------------------------------------------------------------
+
+# pylint: enable=redefined-builtin
+# pylint: enable=arguments-differ
+# pylint: enable=protected-access
+# pylint: enable=line-too-long
+# pylint: enable=global-statement
+# pylint: enable=missing-class-docstring
+# pylint: enable=missing-function-docstring

third_party/stylegan2_official_ops/conv2d_resample.py

@@ -0,0 +1,168 @@
+# python3.7
+
+# 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
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+"""2D convolution with optional up/downsampling.
+
+Please refer to https://github.com/NVlabs/stylegan2-ada-pytorch
+"""
+
+# pylint: disable=line-too-long
+
+import torch
+
+from . import misc
+from . import conv2d_gradfix
+from . import upfirdn2d
+from .upfirdn2d import _parse_padding
+from .upfirdn2d import _get_filter_size
+
+#----------------------------------------------------------------------------
+
+def _get_weight_shape(w):
+    with misc.suppress_tracer_warnings(): # this value will be treated as a constant
+        shape = [int(sz) for sz in w.shape]
+    misc.assert_shape(w, shape)
+    return shape
+
+#----------------------------------------------------------------------------
+
+def _conv2d_wrapper(x, w, stride=1, padding=0, groups=1, transpose=False, flip_weight=True, impl='cuda'):
+    """Wrapper for the underlying `conv2d()` and `conv_transpose2d()` implementations.
+    """
+    out_channels, in_channels_per_group, kh, kw = _get_weight_shape(w)
+
+    # Flip weight if requested.
+    if not flip_weight: # conv2d() actually performs correlation (flip_weight=True) not convolution (flip_weight=False).
+        w = w.flip([2, 3])
+
+    # Workaround performance pitfall in cuDNN 8.0.5, triggered when using
+    # 1x1 kernel + memory_format=channels_last + less than 64 channels.
+    if kw == 1 and kh == 1 and stride == 1 and padding in [0, [0, 0], (0, 0)] and not transpose:
+        if x.stride()[1] == 1 and min(out_channels, in_channels_per_group) < 64:
+            if out_channels <= 4 and groups == 1:
+                in_shape = x.shape
+                x = w.squeeze(3).squeeze(2) @ x.reshape([in_shape[0], in_channels_per_group, -1])
+                x = x.reshape([in_shape[0], out_channels, in_shape[2], in_shape[3]])
+            else:
+                x = x.to(memory_format=torch.contiguous_format)
+                w = w.to(memory_format=torch.contiguous_format)
+                x = conv2d_gradfix.conv2d(x, w, groups=groups, impl=impl)
+            return x.to(memory_format=torch.channels_last)
+
+    # Otherwise => execute using conv2d_gradfix.
+    op = conv2d_gradfix.conv_transpose2d if transpose else conv2d_gradfix.conv2d
+    return op(x, w, stride=stride, padding=padding, groups=groups, impl=impl)
+
+#----------------------------------------------------------------------------
+
+@misc.profiled_function
+def conv2d_resample(x, w, f=None, up=1, down=1, padding=0, groups=1, flip_weight=True, flip_filter=False, impl='cuda'):
+    r"""2D convolution with optional up/downsampling.
+
+    Padding is performed only once at the beginning, not between the operations.
+
+    Args:
+        x:              Input tensor of shape
+                        `[batch_size, in_channels, in_height, in_width]`.
+        w:              Weight tensor of shape
+                        `[out_channels, in_channels//groups, kernel_height, kernel_width]`.
+        f:              Low-pass filter for up/downsampling. Must be prepared beforehand by
+                        calling upfirdn2d.setup_filter(). None = identity (default).
+        up:             Integer upsampling factor (default: 1).
+        down:           Integer downsampling factor (default: 1).
+        padding:        Padding with respect to the upsampled image. Can be a single number
+                        or a list/tuple `[x, y]` or `[x_before, x_after, y_before, y_after]`
+                        (default: 0).
+        groups:         Split input channels into N groups (default: 1).
+        flip_weight:    False = convolution, True = correlation (default: True).
+        flip_filter:    False = convolution, True = correlation (default: False).
+        impl:           Implementation mode of customized ops. 'ref' for native PyTorch
+                        implementation, 'cuda' for `.cu` implementation
+                        (default: 'cuda').
+
+    Returns:
+        Tensor of the shape `[batch_size, num_channels, out_height, out_width]`.
+    """
+    # Validate arguments.
+    assert isinstance(x, torch.Tensor) and (x.ndim == 4)
+    assert isinstance(w, torch.Tensor) and (w.ndim == 4) and (w.dtype == x.dtype)
+    assert f is None or (isinstance(f, torch.Tensor) and f.ndim in [1, 2] and f.dtype == torch.float32)
+    assert isinstance(up, int) and (up >= 1)
+    assert isinstance(down, int) and (down >= 1)
+    assert isinstance(groups, int) and (groups >= 1)
+    out_channels, in_channels_per_group, kh, kw = _get_weight_shape(w)
+    fw, fh = _get_filter_size(f)
+    px0, px1, py0, py1 = _parse_padding(padding)
+
+    # Adjust padding to account for up/downsampling.
+    if up > 1:
+        px0 += (fw + up - 1) // 2
+        px1 += (fw - up) // 2
+        py0 += (fh + up - 1) // 2
+        py1 += (fh - up) // 2
+    if down > 1:
+        px0 += (fw - down + 1) // 2
+        px1 += (fw - down) // 2
+        py0 += (fh - down + 1) // 2
+        py1 += (fh - down) // 2
+
+    # Fast path: 1x1 convolution with downsampling only => downsample first, then convolve.
+    if kw == 1 and kh == 1 and (down > 1 and up == 1):
+        x = upfirdn2d.upfirdn2d(x=x, f=f, down=down, padding=[px0,px1,py0,py1], flip_filter=flip_filter, impl=impl)
+        x = _conv2d_wrapper(x=x, w=w, groups=groups, flip_weight=flip_weight, impl=impl)
+        return x
+
+    # Fast path: 1x1 convolution with upsampling only => convolve first, then upsample.
+    if kw == 1 and kh == 1 and (up > 1 and down == 1):
+        x = _conv2d_wrapper(x=x, w=w, groups=groups, flip_weight=flip_weight, impl=impl)
+        x = upfirdn2d.upfirdn2d(x=x, f=f, up=up, padding=[px0,px1,py0,py1], gain=up**2, flip_filter=flip_filter, impl=impl)
+        return x
+
+    # Fast path: downsampling only => use strided convolution.
+    if down > 1 and up == 1:
+        x = upfirdn2d.upfirdn2d(x=x, f=f, padding=[px0,px1,py0,py1], flip_filter=flip_filter, impl=impl)
+        x = _conv2d_wrapper(x=x, w=w, stride=down, groups=groups, flip_weight=flip_weight, impl=impl)
+        return x
+
+    # Fast path: upsampling with optional downsampling => use transpose strided convolution.
+    if up > 1:
+        if groups == 1:
+            w = w.transpose(0, 1)
+        else:
+            w = w.reshape(groups, out_channels // groups, in_channels_per_group, kh, kw)
+            w = w.transpose(1, 2)
+            w = w.reshape(groups * in_channels_per_group, out_channels // groups, kh, kw)
+        px0 -= kw - 1
+        px1 -= kw - up
+        py0 -= kh - 1
+        py1 -= kh - up
+        pxt = max(min(-px0, -px1), 0)
+        pyt = max(min(-py0, -py1), 0)
+        x = _conv2d_wrapper(x=x, w=w, stride=up, padding=[pyt,pxt], groups=groups, transpose=True, flip_weight=(not flip_weight), impl=impl)
+        x = upfirdn2d.upfirdn2d(x=x, f=f, padding=[px0+pxt,px1+pxt,py0+pyt,py1+pyt], gain=up**2, flip_filter=flip_filter, impl=impl)
+        if down > 1:
+            x = upfirdn2d.upfirdn2d(x=x, f=f, down=down, flip_filter=flip_filter, impl=impl)
+        return x
+
+    # Fast path: no up/downsampling, padding supported by the underlying implementation => use plain conv2d.
+    if up == 1 and down == 1:
+        if px0 == px1 and py0 == py1 and px0 >= 0 and py0 >= 0:
+            return _conv2d_wrapper(x=x, w=w, padding=[py0,px0], groups=groups, flip_weight=flip_weight, impl=impl)
+
+    # Fallback: Generic reference implementation.
+    x = upfirdn2d.upfirdn2d(x=x, f=(f if up > 1 else None), up=up, padding=[px0,px1,py0,py1], gain=up**2, flip_filter=flip_filter, impl=impl)
+    x = _conv2d_wrapper(x=x, w=w, groups=groups, flip_weight=flip_weight, impl=impl)
+    if down > 1:
+        x = upfirdn2d.upfirdn2d(x=x, f=f, down=down, flip_filter=flip_filter, impl=impl)
+    return x
+
+#----------------------------------------------------------------------------
+
+# pylint: enable=line-too-long

third_party/stylegan2_official_ops/custom_ops.py

@@ -0,0 +1,159 @@
+# python3.7
+
+# 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
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+"""Utility functions to setup customized operators.
+
+Please refer to https://github.com/NVlabs/stylegan2-ada-pytorch
+"""
+
+# pylint: disable=line-too-long
+# pylint: disable=missing-function-docstring
+# pylint: disable=useless-suppression
+# pylint: disable=inconsistent-quotes
+
+import os
+import glob
+import importlib
+import hashlib
+import shutil
+from pathlib import Path
+
+import torch
+from torch.utils.file_baton import FileBaton
+import torch.utils.cpp_extension
+
+#----------------------------------------------------------------------------
+# Global options.
+
+verbosity = 'none' # Verbosity level: 'none', 'brief', 'full'
+
+#----------------------------------------------------------------------------
+# Internal helper funcs.
+
+def _find_compiler_bindir():
+    patterns = [
+        'C:/Program Files (x86)/Microsoft Visual Studio/*/Professional/VC/Tools/MSVC/*/bin/Hostx64/x64',
+        'C:/Program Files (x86)/Microsoft Visual Studio/*/BuildTools/VC/Tools/MSVC/*/bin/Hostx64/x64',
+        'C:/Program Files (x86)/Microsoft Visual Studio/*/Community/VC/Tools/MSVC/*/bin/Hostx64/x64',
+        'C:/Program Files (x86)/Microsoft Visual Studio */vc/bin',
+    ]
+    for pattern in patterns:
+        matches = sorted(glob.glob(pattern))
+        if len(matches):
+            return matches[-1]
+    return None
+
+def _find_compiler_bindir_posix():
+    patterns = [
+        '/usr/local/cuda/bin'
+    ]
+    for pattern in patterns:
+        matches = sorted(glob.glob(pattern))
+        if len(matches):
+            return matches[-1]
+    return None
+
+#----------------------------------------------------------------------------
+# Main entry point for compiling and loading C++/CUDA plugins.
+
+_cached_plugins = dict()
+
+def get_plugin(module_name, sources, **build_kwargs):
+    assert verbosity in ['none', 'brief', 'full']
+
+    # Already cached?
+    if module_name in _cached_plugins:
+        return _cached_plugins[module_name]
+
+    # Print status.
+    if verbosity == 'full':
+        print(f'Setting up PyTorch plugin "{module_name}"...')
+    elif verbosity == 'brief':
+        print(f'Setting up PyTorch plugin "{module_name}"... ', end='', flush=True)
+
+    try: # pylint: disable=too-many-nested-blocks
+        # Make sure we can find the necessary compiler binaries.
+        if os.name == 'nt' and os.system("where cl.exe >nul 2>nul") != 0:
+            compiler_bindir = _find_compiler_bindir()
+            if compiler_bindir is None:
+                raise RuntimeError(f'Could not find MSVC/GCC/CLANG installation on this computer. Check _find_compiler_bindir() in "{__file__}".')
+            os.environ['PATH'] += ';' + compiler_bindir
+
+        elif os.name == 'posix':
+            compiler_bindir = _find_compiler_bindir_posix()
+            if compiler_bindir is None:
+                raise RuntimeError(f'Could not find NVCC installation on this computer. Check _find_compiler_bindir_posix() in "{__file__}".')
+            os.environ['PATH'] += ';' + compiler_bindir
+
+        # Compile and load.
+        verbose_build = (verbosity == 'full')
+
+        # Incremental build md5sum trickery.  Copies all the input source files
+        # into a cached build directory under a combined md5 digest of the input
+        # source files.  Copying is done only if the combined digest has changed.
+        # This keeps input file timestamps and filenames the same as in previous
+        # extension builds, allowing for fast incremental rebuilds.
+        #
+        # This optimization is done only in case all the source files reside in
+        # a single directory (just for simplicity) and if the TORCH_EXTENSIONS_DIR
+        # environment variable is set (we take this as a signal that the user
+        # actually cares about this.)
+        source_dirs_set = set(os.path.dirname(source) for source in sources)
+        if len(source_dirs_set) == 1 and ('TORCH_EXTENSIONS_DIR' in os.environ):
+            all_source_files = sorted(list(x for x in Path(list(source_dirs_set)[0]).iterdir() if x.is_file()))
+
+            # Compute a combined hash digest for all source files in the same
+            # custom op directory (usually .cu, .cpp, .py and .h files).
+            hash_md5 = hashlib.md5()
+            for src in all_source_files:
+                with open(src, 'rb') as f:
+                    hash_md5.update(f.read())
+            build_dir = torch.utils.cpp_extension._get_build_directory(module_name, verbose=verbose_build) # pylint: disable=protected-access
+            digest_build_dir = os.path.join(build_dir, hash_md5.hexdigest())
+
+            if not os.path.isdir(digest_build_dir):
+                os.makedirs(digest_build_dir, exist_ok=True)
+                baton = FileBaton(os.path.join(digest_build_dir, 'lock'))
+                if baton.try_acquire():
+                    try:
+                        for src in all_source_files:
+                            shutil.copyfile(src, os.path.join(digest_build_dir, os.path.basename(src)))
+                    finally:
+                        baton.release()
+                else:
+                    # Someone else is copying source files under the digest dir,
+                    # wait until done and continue.
+                    baton.wait()
+            digest_sources = [os.path.join(digest_build_dir, os.path.basename(x)) for x in sources]
+            torch.utils.cpp_extension.load(name=module_name, build_directory=build_dir,
+                verbose=verbose_build, sources=digest_sources, **build_kwargs)
+        else:
+            torch.utils.cpp_extension.load(name=module_name, verbose=verbose_build, sources=sources, **build_kwargs)
+        module = importlib.import_module(module_name)
+
+    except:
+        if verbosity == 'brief':
+            print('Failed!')
+        raise
+
+    # Print status and add to cache.
+    if verbosity == 'full':
+        print(f'Done setting up PyTorch plugin "{module_name}".')
+    elif verbosity == 'brief':
+        print('Done.')
+    _cached_plugins[module_name] = module
+    return module
+
+#----------------------------------------------------------------------------
+
+# pylint: enable=line-too-long
+# pylint: enable=missing-function-docstring
+# pylint: enable=useless-suppression
+# pylint: enable=inconsistent-quotes

third_party/stylegan2_official_ops/fma.py

@@ -0,0 +1,73 @@
+# python3.7
+
+# 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
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+"""Fused multiply-add, with slightly faster gradients than `torch.addcmul()`.
+
+Please refer to https://github.com/NVlabs/stylegan2-ada-pytorch
+"""
+
+# pylint: disable=line-too-long
+# pylint: disable=missing-function-docstring
+
+import torch
+
+#----------------------------------------------------------------------------
+
+def fma(a, b, c, impl='cuda'): # => a * b + c
+    if impl == 'cuda':
+        return _FusedMultiplyAdd.apply(a, b, c)
+    return torch.addcmul(c, a, b)
+
+#----------------------------------------------------------------------------
+
+class _FusedMultiplyAdd(torch.autograd.Function): # a * b + c
+    @staticmethod
+    def forward(ctx, a, b, c): # pylint: disable=arguments-differ
+        out = torch.addcmul(c, a, b)
+        ctx.save_for_backward(a, b)
+        ctx.c_shape = c.shape
+        return out
+
+    @staticmethod
+    def backward(ctx, dout): # pylint: disable=arguments-differ
+        a, b = ctx.saved_tensors
+        c_shape = ctx.c_shape
+        da = None
+        db = None
+        dc = None
+
+        if ctx.needs_input_grad[0]:
+            da = _unbroadcast(dout * b, a.shape)
+
+        if ctx.needs_input_grad[1]:
+            db = _unbroadcast(dout * a, b.shape)
+
+        if ctx.needs_input_grad[2]:
+            dc = _unbroadcast(dout, c_shape)
+
+        return da, db, dc
+
+#----------------------------------------------------------------------------
+
+def _unbroadcast(x, shape):
+    extra_dims = x.ndim - len(shape)
+    assert extra_dims >= 0
+    dim = [i for i in range(x.ndim) if x.shape[i] > 1 and (i < extra_dims or shape[i - extra_dims] == 1)]
+    if len(dim):
+        x = x.sum(dim=dim, keepdim=True)
+    if extra_dims:
+        x = x.reshape(-1, *x.shape[extra_dims+1:])
+    assert x.shape == shape
+    return x
+
+#----------------------------------------------------------------------------
+
+# pylint: enable=line-too-long
+# pylint: enable=missing-function-docstring

third_party/stylegan2_official_ops/grid_sample_gradfix.py

@@ -0,0 +1,98 @@
+# python3.7
+
+# 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
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+"""Custom replacement for `torch.nn.functional.grid_sample`.
+
+This is useful for differentiable augmentation. This customized operator
+supports arbitrarily high order gradients between the input and output. Only
+works on 2D images and assumes `mode=bilinear`, `padding_mode=zeros`, and
+`align_corners=False`.
+
+Please refer to https://github.com/NVlabs/stylegan2-ada-pytorch
+"""
+
+# pylint: disable=redefined-builtin
+# pylint: disable=arguments-differ
+# pylint: disable=protected-access
+# pylint: disable=line-too-long
+# pylint: disable=missing-function-docstring
+
+import warnings
+import torch
+
+#----------------------------------------------------------------------------
+
+enabled = True  # Enable the custom op by setting this to true.
+
+#----------------------------------------------------------------------------
+
+def grid_sample(input, grid, impl='cuda'):
+    if impl == 'cuda' and _should_use_custom_op():
+        return _GridSample2dForward.apply(input, grid)
+    return torch.nn.functional.grid_sample(input=input, grid=grid, mode='bilinear', padding_mode='zeros', align_corners=False)
+
+#----------------------------------------------------------------------------
+
+def _should_use_custom_op():
+    if not enabled:
+        return False
+    if any(torch.__version__.startswith(x) for x in ['1.7.', '1.8.', '1.9']):
+        return True
+    warnings.warn(f'grid_sample_gradfix not supported on PyTorch {torch.__version__}. Falling back to torch.nn.functional.grid_sample().')
+    return False
+
+#----------------------------------------------------------------------------
+
+class _GridSample2dForward(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input, grid):
+        assert input.ndim == 4
+        assert grid.ndim == 4
+        output = torch.nn.functional.grid_sample(input=input, grid=grid, mode='bilinear', padding_mode='zeros', align_corners=False)
+        ctx.save_for_backward(input, grid)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        input, grid = ctx.saved_tensors
+        grad_input, grad_grid = _GridSample2dBackward.apply(grad_output, input, grid)
+        return grad_input, grad_grid
+
+#----------------------------------------------------------------------------
+
+class _GridSample2dBackward(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, grad_output, input, grid):
+        op = torch._C._jit_get_operation('aten::grid_sampler_2d_backward')
+        grad_input, grad_grid = op(grad_output, input, grid, 0, 0, False)
+        ctx.save_for_backward(grid)
+        return grad_input, grad_grid
+
+    @staticmethod
+    def backward(ctx, grad2_grad_input, grad2_grad_grid):
+        _ = grad2_grad_grid # unused
+        grid, = ctx.saved_tensors
+        grad2_grad_output = None
+        grad2_input = None
+        grad2_grid = None
+
+        if ctx.needs_input_grad[0]:
+            grad2_grad_output = _GridSample2dForward.apply(grad2_grad_input, grid)
+
+        assert not ctx.needs_input_grad[2]
+        return grad2_grad_output, grad2_input, grad2_grid
+
+#----------------------------------------------------------------------------
+
+# pylint: enable=redefined-builtin
+# pylint: enable=arguments-differ
+# pylint: enable=protected-access
+# pylint: enable=line-too-long
+# pylint: enable=missing-function-docstring

third_party/stylegan2_official_ops/misc.py

@@ -0,0 +1,281 @@
+# python3.7
+
+# 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
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+"""Misc functions for customized operations.
+
+Please refer to https://github.com/NVlabs/stylegan2-ada-pytorch
+"""
+
+# pylint: disable=line-too-long
+# pylint: disable=missing-class-docstring
+# pylint: disable=missing-function-docstring
+# pylint: disable=use-maxsplit-arg
+# pylint: disable=unnecessary-comprehension
+
+import re
+import contextlib
+import warnings
+from easydict import EasyDict
+import numpy as np
+import torch
+
+#----------------------------------------------------------------------------
+# Cached construction of constant tensors. Avoids CPU=>GPU copy when the
+# same constant is used multiple times.
+
+_constant_cache = dict()
+
+def constant(value, shape=None, dtype=None, device=None, memory_format=None):
+    value = np.asarray(value)
+    if shape is not None:
+        shape = tuple(shape)
+    if dtype is None:
+        dtype = torch.get_default_dtype()
+    if device is None:
+        device = torch.device('cpu')
+    if memory_format is None:
+        memory_format = torch.contiguous_format
+
+    key = (value.shape, value.dtype, value.tobytes(), shape, dtype, device, memory_format)
+    tensor = _constant_cache.get(key, None)
+    if tensor is None:
+        tensor = torch.as_tensor(value.copy(), dtype=dtype, device=device)
+        if shape is not None:
+            tensor, _ = torch.broadcast_tensors(tensor, torch.empty(shape))
+        tensor = tensor.contiguous(memory_format=memory_format)
+        _constant_cache[key] = tensor
+    return tensor
+
+#----------------------------------------------------------------------------
+# Replace NaN/Inf with specified numerical values.
+
+try:
+    nan_to_num = torch.nan_to_num # 1.8.0a0
+except AttributeError:
+    def nan_to_num(input, nan=0.0, posinf=None, neginf=None, *, out=None): # pylint: disable=redefined-builtin
+        assert isinstance(input, torch.Tensor)
+        if posinf is None:
+            posinf = torch.finfo(input.dtype).max
+        if neginf is None:
+            neginf = torch.finfo(input.dtype).min
+        assert nan == 0
+        return torch.clamp(input.unsqueeze(0).nansum(0), min=neginf, max=posinf, out=out)
+
+#----------------------------------------------------------------------------
+# Symbolic assert.
+
+try:
+    symbolic_assert = torch._assert # 1.8.0a0 # pylint: disable=protected-access
+except AttributeError:
+    symbolic_assert = torch.Assert # 1.7.0
+
+#----------------------------------------------------------------------------
+# Context manager to suppress known warnings in torch.jit.trace().
+
+class suppress_tracer_warnings(warnings.catch_warnings):
+    def __enter__(self):
+        super().__enter__()
+        warnings.simplefilter('ignore', category=torch.jit.TracerWarning)
+        return self
+
+#----------------------------------------------------------------------------
+# Assert that the shape of a tensor matches the given list of integers.
+# None indicates that the size of a dimension is allowed to vary.
+# Performs symbolic assertion when used in torch.jit.trace().
+
+def assert_shape(tensor, ref_shape):
+    if tensor.ndim != len(ref_shape):
+        raise AssertionError(f'Wrong number of dimensions: got {tensor.ndim}, expected {len(ref_shape)}')
+    for idx, (size, ref_size) in enumerate(zip(tensor.shape, ref_shape)):
+        if ref_size is None:
+            pass
+        elif isinstance(ref_size, torch.Tensor):
+            with suppress_tracer_warnings(): # as_tensor results are registered as constants
+                symbolic_assert(torch.equal(torch.as_tensor(size), ref_size), f'Wrong size for dimension {idx}')
+        elif isinstance(size, torch.Tensor):
+            with suppress_tracer_warnings(): # as_tensor results are registered as constants
+                symbolic_assert(torch.equal(size, torch.as_tensor(ref_size)), f'Wrong size for dimension {idx}: expected {ref_size}')
+        elif size != ref_size:
+            raise AssertionError(f'Wrong size for dimension {idx}: got {size}, expected {ref_size}')
+
+#----------------------------------------------------------------------------
+# Function decorator that calls torch.autograd.profiler.record_function().
+
+def profiled_function(fn):
+    def decorator(*args, **kwargs):
+        with torch.autograd.profiler.record_function(fn.__name__):
+            return fn(*args, **kwargs)
+    decorator.__name__ = fn.__name__
+    return decorator
+
+#----------------------------------------------------------------------------
+# Sampler for torch.utils.data.DataLoader that loops over the dataset
+# indefinitely, shuffling items as it goes.
+
+class InfiniteSampler(torch.utils.data.Sampler):
+    def __init__(self, dataset, rank=0, num_replicas=1, shuffle=True, seed=0, window_size=0.5):
+        assert len(dataset) > 0
+        assert num_replicas > 0
+        assert 0 <= rank < num_replicas
+        assert 0 <= window_size <= 1
+        super().__init__(dataset)
+        self.dataset = dataset
+        self.rank = rank
+        self.num_replicas = num_replicas
+        self.shuffle = shuffle
+        self.seed = seed
+        self.window_size = window_size
+
+    def __iter__(self):
+        order = np.arange(len(self.dataset))
+        rnd = None
+        window = 0
+        if self.shuffle:
+            rnd = np.random.RandomState(self.seed)
+            rnd.shuffle(order)
+            window = int(np.rint(order.size * self.window_size))
+
+        idx = 0
+        while True:
+            i = idx % order.size
+            if idx % self.num_replicas == self.rank:
+                yield order[i]
+            if window >= 2:
+                j = (i - rnd.randint(window)) % order.size
+                order[i], order[j] = order[j], order[i]
+            idx += 1
+
+#----------------------------------------------------------------------------
+# Utilities for operating with torch.nn.Module parameters and buffers.
+
+def params_and_buffers(module):
+    assert isinstance(module, torch.nn.Module)
+    return list(module.parameters()) + list(module.buffers())
+
+def named_params_and_buffers(module):
+    assert isinstance(module, torch.nn.Module)
+    return list(module.named_parameters()) + list(module.named_buffers())
+
+def copy_params_and_buffers(src_module, dst_module, require_all=False):
+    assert isinstance(src_module, torch.nn.Module)
+    assert isinstance(dst_module, torch.nn.Module)
+    src_tensors = {name: tensor for name, tensor in named_params_and_buffers(src_module)}
+    for name, tensor in named_params_and_buffers(dst_module):
+        assert (name in src_tensors) or (not require_all)
+        if name in src_tensors:
+            tensor.copy_(src_tensors[name].detach()).requires_grad_(tensor.requires_grad)
+
+#----------------------------------------------------------------------------
+# Context manager for easily enabling/disabling DistributedDataParallel
+# synchronization.
+
+@contextlib.contextmanager
+def ddp_sync(module, sync):
+    assert isinstance(module, torch.nn.Module)
+    if sync or not isinstance(module, torch.nn.parallel.DistributedDataParallel):
+        yield
+    else:
+        with module.no_sync():
+            yield
+
+#----------------------------------------------------------------------------
+# Check DistributedDataParallel consistency across processes.
+
+def check_ddp_consistency(module, ignore_regex=None):
+    assert isinstance(module, torch.nn.Module)
+    for name, tensor in named_params_and_buffers(module):
+        fullname = type(module).__name__ + '.' + name
+        if ignore_regex is not None and re.fullmatch(ignore_regex, fullname):
+            continue
+        tensor = tensor.detach()
+        other = tensor.clone()
+        torch.distributed.broadcast(tensor=other, src=0)
+        assert (nan_to_num(tensor) == nan_to_num(other)).all(), fullname
+
+#----------------------------------------------------------------------------
+# Print summary table of module hierarchy.
+
+def print_module_summary(module, inputs, max_nesting=3, skip_redundant=True):
+    assert isinstance(module, torch.nn.Module)
+    assert not isinstance(module, torch.jit.ScriptModule)
+    assert isinstance(inputs, (tuple, list))
+
+    # Register hooks.
+    entries = []
+    nesting = [0]
+    def pre_hook(_mod, _inputs):
+        nesting[0] += 1
+    def post_hook(mod, _inputs, outputs):
+        nesting[0] -= 1
+        if nesting[0] <= max_nesting:
+            outputs = list(outputs) if isinstance(outputs, (tuple, list)) else [outputs]
+            outputs = [t for t in outputs if isinstance(t, torch.Tensor)]
+            entries.append(EasyDict(mod=mod, outputs=outputs))
+    hooks = [mod.register_forward_pre_hook(pre_hook) for mod in module.modules()]
+    hooks += [mod.register_forward_hook(post_hook) for mod in module.modules()]
+
+    # Run module.
+    outputs = module(*inputs)
+    for hook in hooks:
+        hook.remove()
+
+    # Identify unique outputs, parameters, and buffers.
+    tensors_seen = set()
+    for e in entries:
+        e.unique_params = [t for t in e.mod.parameters() if id(t) not in tensors_seen]
+        e.unique_buffers = [t for t in e.mod.buffers() if id(t) not in tensors_seen]
+        e.unique_outputs = [t for t in e.outputs if id(t) not in tensors_seen]
+        tensors_seen |= {id(t) for t in e.unique_params + e.unique_buffers + e.unique_outputs}
+
+    # Filter out redundant entries.
+    if skip_redundant:
+        entries = [e for e in entries if len(e.unique_params) or len(e.unique_buffers) or len(e.unique_outputs)]
+
+    # Construct table.
+    rows = [[type(module).__name__, 'Parameters', 'Buffers', 'Output shape', 'Datatype']]
+    rows += [['---'] * len(rows[0])]
+    param_total = 0
+    buffer_total = 0
+    submodule_names = {mod: name for name, mod in module.named_modules()}
+    for e in entries:
+        name = '<top-level>' if e.mod is module else submodule_names[e.mod]
+        param_size = sum(t.numel() for t in e.unique_params)
+        buffer_size = sum(t.numel() for t in e.unique_buffers)
+        output_shapes = [str(list(e.outputs[0].shape)) for t in e.outputs]
+        output_dtypes = [str(t.dtype).split('.')[-1] for t in e.outputs]
+        rows += [[
+            name + (':0' if len(e.outputs) >= 2 else ''),
+            str(param_size) if param_size else '-',
+            str(buffer_size) if buffer_size else '-',
+            (output_shapes + ['-'])[0],
+            (output_dtypes + ['-'])[0],
+        ]]
+        for idx in range(1, len(e.outputs)):
+            rows += [[name + f':{idx}', '-', '-', output_shapes[idx], output_dtypes[idx]]]
+        param_total += param_size
+        buffer_total += buffer_size
+    rows += [['---'] * len(rows[0])]
+    rows += [['Total', str(param_total), str(buffer_total), '-', '-']]
+
+    # Print table.
+    widths = [max(len(cell) for cell in column) for column in zip(*rows)]
+    print()
+    for row in rows:
+        print('  '.join(cell + ' ' * (width - len(cell)) for cell, width in zip(row, widths)))
+    print()
+    return outputs
+
+#----------------------------------------------------------------------------
+
+# pylint: enable=line-too-long
+# pylint: enable=missing-class-docstring
+# pylint: enable=missing-function-docstring
+# pylint: enable=use-maxsplit-arg
+# pylint: enable=unnecessary-comprehension

third_party/stylegan2_official_ops/upfirdn2d.cpp

@@ -0,0 +1,103 @@
+// 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
+// and any modifications thereto.  Any use, reproduction, disclosure or
+// distribution of this software and related documentation without an express
+// license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+#include <torch/extension.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include "upfirdn2d.h"
+
+//------------------------------------------------------------------------
+
+static torch::Tensor upfirdn2d(torch::Tensor x, torch::Tensor f, int upx, int upy, int downx, int downy, int padx0, int padx1, int pady0, int pady1, bool flip, float gain)
+{
+    // Validate arguments.
+    TORCH_CHECK(x.is_cuda(), "x must reside on CUDA device");
+    TORCH_CHECK(f.device() == x.device(), "f must reside on the same device as x");
+    TORCH_CHECK(f.dtype() == torch::kFloat, "f must be float32");
+    TORCH_CHECK(x.numel() <= INT_MAX, "x is too large");
+    TORCH_CHECK(f.numel() <= INT_MAX, "f is too large");
+    TORCH_CHECK(x.dim() == 4, "x must be rank 4");
+    TORCH_CHECK(f.dim() == 2, "f must be rank 2");
+    TORCH_CHECK(f.size(0) >= 1 && f.size(1) >= 1, "f must be at least 1x1");
+    TORCH_CHECK(upx >= 1 && upy >= 1, "upsampling factor must be at least 1");
+    TORCH_CHECK(downx >= 1 && downy >= 1, "downsampling factor must be at least 1");
+
+    // Create output tensor.
+    const at::cuda::OptionalCUDAGuard device_guard(device_of(x));
+    int outW = ((int)x.size(3) * upx + padx0 + padx1 - (int)f.size(1) + downx) / downx;
+    int outH = ((int)x.size(2) * upy + pady0 + pady1 - (int)f.size(0) + downy) / downy;
+    TORCH_CHECK(outW >= 1 && outH >= 1, "output must be at least 1x1");
+    torch::Tensor y = torch::empty({x.size(0), x.size(1), outH, outW}, x.options(), x.suggest_memory_format());
+    TORCH_CHECK(y.numel() <= INT_MAX, "output is too large");
+
+    // Initialize CUDA kernel parameters.
+    upfirdn2d_kernel_params p;
+    p.x             = x.data_ptr();
+    p.f             = f.data_ptr<float>();
+    p.y             = y.data_ptr();
+    p.up            = make_int2(upx, upy);
+    p.down          = make_int2(downx, downy);
+    p.pad0          = make_int2(padx0, pady0);
+    p.flip          = (flip) ? 1 : 0;
+    p.gain          = gain;
+    p.inSize        = make_int4((int)x.size(3), (int)x.size(2), (int)x.size(1), (int)x.size(0));
+    p.inStride      = make_int4((int)x.stride(3), (int)x.stride(2), (int)x.stride(1), (int)x.stride(0));
+    p.filterSize    = make_int2((int)f.size(1), (int)f.size(0));
+    p.filterStride  = make_int2((int)f.stride(1), (int)f.stride(0));
+    p.outSize       = make_int4((int)y.size(3), (int)y.size(2), (int)y.size(1), (int)y.size(0));
+    p.outStride     = make_int4((int)y.stride(3), (int)y.stride(2), (int)y.stride(1), (int)y.stride(0));
+    p.sizeMajor     = (p.inStride.z == 1) ? p.inSize.w : p.inSize.w * p.inSize.z;
+    p.sizeMinor     = (p.inStride.z == 1) ? p.inSize.z : 1;
+
+    // Choose CUDA kernel.
+    upfirdn2d_kernel_spec spec;
+    AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), "upfirdn2d_cuda", [&]
+    {
+        spec = choose_upfirdn2d_kernel<scalar_t>(p);
+    });
+
+    // Set looping options.
+    p.loopMajor     = (p.sizeMajor - 1) / 16384 + 1;
+    p.loopMinor     = spec.loopMinor;
+    p.loopX         = spec.loopX;
+    p.launchMinor   = (p.sizeMinor - 1) / p.loopMinor + 1;
+    p.launchMajor   = (p.sizeMajor - 1) / p.loopMajor + 1;
+
+    // Compute grid size.
+    dim3 blockSize, gridSize;
+    if (spec.tileOutW < 0) // large
+    {
+        blockSize = dim3(4, 32, 1);
+        gridSize = dim3(
+            ((p.outSize.y - 1) / blockSize.x + 1) * p.launchMinor,
+            (p.outSize.x - 1) / (blockSize.y * p.loopX) + 1,
+            p.launchMajor);
+    }
+    else // small
+    {
+        blockSize = dim3(256, 1, 1);
+        gridSize = dim3(
+            ((p.outSize.y - 1) / spec.tileOutH + 1) * p.launchMinor,
+            (p.outSize.x - 1) / (spec.tileOutW * p.loopX) + 1,
+            p.launchMajor);
+    }
+
+    // Launch CUDA kernel.
+    void* args[] = {&p};
+    AT_CUDA_CHECK(cudaLaunchKernel(spec.kernel, gridSize, blockSize, args, 0, at::cuda::getCurrentCUDAStream()));
+    return y;
+}
+
+//------------------------------------------------------------------------
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
+{
+    m.def("upfirdn2d", &upfirdn2d);
+}
+
+//------------------------------------------------------------------------

third_party/stylegan2_official_ops/upfirdn2d.cu

@@ -0,0 +1,350 @@
+// 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
+// and any modifications thereto.  Any use, reproduction, disclosure or
+// distribution of this software and related documentation without an express
+// license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+#include <c10/util/Half.h>
+#include "upfirdn2d.h"
+
+//------------------------------------------------------------------------
+// Helpers.
+
+template <class T> struct InternalType;
+template <> struct InternalType<double>     { typedef double scalar_t; };
+template <> struct InternalType<float>      { typedef float  scalar_t; };
+template <> struct InternalType<c10::Half>  { typedef float  scalar_t; };
+
+static __device__ __forceinline__ int floor_div(int a, int b)
+{
+    int t = 1 - a / b;
+    return (a + t * b) / b - t;
+}
+
+//------------------------------------------------------------------------
+// Generic CUDA implementation for large filters.
+
+template <class T> static __global__ void upfirdn2d_kernel_large(upfirdn2d_kernel_params p)
+{
+    typedef typename InternalType<T>::scalar_t scalar_t;
+
+    // Calculate thread index.
+    int minorBase = blockIdx.x * blockDim.x + threadIdx.x;
+    int outY = minorBase / p.launchMinor;
+    minorBase -= outY * p.launchMinor;
+    int outXBase = blockIdx.y * p.loopX * blockDim.y + threadIdx.y;
+    int majorBase = blockIdx.z * p.loopMajor;
+    if (outXBase >= p.outSize.x | outY >= p.outSize.y | majorBase >= p.sizeMajor)
+        return;
+
+    // Setup Y receptive field.
+    int midY = outY * p.down.y + p.up.y - 1 - p.pad0.y;
+    int inY = min(max(floor_div(midY, p.up.y), 0), p.inSize.y);
+    int h = min(max(floor_div(midY + p.filterSize.y, p.up.y), 0), p.inSize.y) - inY;
+    int filterY = midY + p.filterSize.y - (inY + 1) * p.up.y;
+    if (p.flip)
+        filterY = p.filterSize.y - 1 - filterY;
+
+    // Loop over major, minor, and X.
+    for (int majorIdx = 0, major = majorBase; majorIdx < p.loopMajor & major < p.sizeMajor; majorIdx++, major++)
+    for (int minorIdx = 0, minor = minorBase; minorIdx < p.loopMinor & minor < p.sizeMinor; minorIdx++, minor += p.launchMinor)
+    {
+        int nc = major * p.sizeMinor + minor;
+        int n = nc / p.inSize.z;
+        int c = nc - n * p.inSize.z;
+        for (int loopX = 0, outX = outXBase; loopX < p.loopX & outX < p.outSize.x; loopX++, outX += blockDim.y)
+        {
+            // Setup X receptive field.
+            int midX = outX * p.down.x + p.up.x - 1 - p.pad0.x;
+            int inX = min(max(floor_div(midX, p.up.x), 0), p.inSize.x);
+            int w = min(max(floor_div(midX + p.filterSize.x, p.up.x), 0), p.inSize.x) - inX;
+            int filterX = midX + p.filterSize.x - (inX + 1) * p.up.x;
+            if (p.flip)
+                filterX = p.filterSize.x - 1 - filterX;
+
+            // Initialize pointers.
+            const T* xp = &((const T*)p.x)[inX * p.inStride.x + inY * p.inStride.y + c * p.inStride.z + n * p.inStride.w];
+            const float* fp = &p.f[filterX * p.filterStride.x + filterY * p.filterStride.y];
+            int filterStepX = ((p.flip) ? p.up.x : -p.up.x) * p.filterStride.x;
+            int filterStepY = ((p.flip) ? p.up.y : -p.up.y) * p.filterStride.y;
+
+            // Inner loop.
+            scalar_t v = 0;
+            for (int y = 0; y < h; y++)
+            {
+                for (int x = 0; x < w; x++)
+                {
+                    v += (scalar_t)(*xp) * (scalar_t)(*fp);
+                    xp += p.inStride.x;
+                    fp += filterStepX;
+                }
+                xp += p.inStride.y - w * p.inStride.x;
+                fp += filterStepY - w * filterStepX;
+            }
+
+            // Store result.
+            v *= p.gain;
+            ((T*)p.y)[outX * p.outStride.x + outY * p.outStride.y + c * p.outStride.z + n * p.outStride.w] = (T)v;
+        }
+    }
+}
+
+//------------------------------------------------------------------------
+// Specialized CUDA implementation for small filters.
+
+template <class T, int upx, int upy, int downx, int downy, int filterW, int filterH, int tileOutW, int tileOutH, int loopMinor>
+static __global__ void upfirdn2d_kernel_small(upfirdn2d_kernel_params p)
+{
+    typedef typename InternalType<T>::scalar_t scalar_t;
+    const int tileInW = ((tileOutW - 1) * downx + filterW - 1) / upx + 1;
+    const int tileInH = ((tileOutH - 1) * downy + filterH - 1) / upy + 1;
+    __shared__ volatile scalar_t sf[filterH][filterW];
+    __shared__ volatile scalar_t sx[tileInH][tileInW][loopMinor];
+
+    // Calculate tile index.
+    int minorBase = blockIdx.x;
+    int tileOutY = minorBase / p.launchMinor;
+    minorBase -= tileOutY * p.launchMinor;
+    minorBase *= loopMinor;
+    tileOutY *= tileOutH;
+    int tileOutXBase = blockIdx.y * p.loopX * tileOutW;
+    int majorBase = blockIdx.z * p.loopMajor;
+    if (tileOutXBase >= p.outSize.x | tileOutY >= p.outSize.y | majorBase >= p.sizeMajor)
+        return;
+
+    // Load filter (flipped).
+    for (int tapIdx = threadIdx.x; tapIdx < filterH * filterW; tapIdx += blockDim.x)
+    {
+        int fy = tapIdx / filterW;
+        int fx = tapIdx - fy * filterW;
+        scalar_t v = 0;
+        if (fx < p.filterSize.x & fy < p.filterSize.y)
+        {
+            int ffx = (p.flip) ? fx : p.filterSize.x - 1 - fx;
+            int ffy = (p.flip) ? fy : p.filterSize.y - 1 - fy;
+            v = (scalar_t)p.f[ffx * p.filterStride.x + ffy * p.filterStride.y];
+        }
+        sf[fy][fx] = v;
+    }
+
+    // Loop over major and X.
+    for (int majorIdx = 0, major = majorBase; majorIdx < p.loopMajor & major < p.sizeMajor; majorIdx++, major++)
+    {
+        int baseNC = major * p.sizeMinor + minorBase;
+        int n = baseNC / p.inSize.z;
+        int baseC = baseNC - n * p.inSize.z;
+        for (int loopX = 0, tileOutX = tileOutXBase; loopX < p.loopX & tileOutX < p.outSize.x; loopX++, tileOutX += tileOutW)
+        {
+            // Load input pixels.
+            int tileMidX = tileOutX * downx + upx - 1 - p.pad0.x;
+            int tileMidY = tileOutY * downy + upy - 1 - p.pad0.y;
+            int tileInX = floor_div(tileMidX, upx);
+            int tileInY = floor_div(tileMidY, upy);
+            __syncthreads();
+            for (int inIdx = threadIdx.x; inIdx < tileInH * tileInW * loopMinor; inIdx += blockDim.x)
+            {
+                int relC = inIdx;
+                int relInX = relC / loopMinor;
+                int relInY = relInX / tileInW;
+                relC -= relInX * loopMinor;
+                relInX -= relInY * tileInW;
+                int c = baseC + relC;
+                int inX = tileInX + relInX;
+                int inY = tileInY + relInY;
+                scalar_t v = 0;
+                if (inX >= 0 & inY >= 0 & inX < p.inSize.x & inY < p.inSize.y & c < p.inSize.z)
+                    v = (scalar_t)((const T*)p.x)[inX * p.inStride.x + inY * p.inStride.y + c * p.inStride.z + n * p.inStride.w];
+                sx[relInY][relInX][relC] = v;
+            }
+
+            // Loop over output pixels.
+            __syncthreads();
+            for (int outIdx = threadIdx.x; outIdx < tileOutH * tileOutW * loopMinor; outIdx += blockDim.x)
+            {
+                int relC = outIdx;
+                int relOutX = relC / loopMinor;
+                int relOutY = relOutX / tileOutW;
+                relC -= relOutX * loopMinor;
+                relOutX -= relOutY * tileOutW;
+                int c = baseC + relC;
+                int outX = tileOutX + relOutX;
+                int outY = tileOutY + relOutY;
+
+                // Setup receptive field.
+                int midX = tileMidX + relOutX * downx;
+                int midY = tileMidY + relOutY * downy;
+                int inX = floor_div(midX, upx);
+                int inY = floor_div(midY, upy);
+                int relInX = inX - tileInX;
+                int relInY = inY - tileInY;
+                int filterX = (inX + 1) * upx - midX - 1; // flipped
+                int filterY = (inY + 1) * upy - midY - 1; // flipped
+
+                // Inner loop.
+                if (outX < p.outSize.x & outY < p.outSize.y & c < p.outSize.z)
+                {
+                    scalar_t v = 0;
+                    #pragma unroll
+                    for (int y = 0; y < filterH / upy; y++)
+                        #pragma unroll
+                        for (int x = 0; x < filterW / upx; x++)
+                            v += sx[relInY + y][relInX + x][relC] * sf[filterY + y * upy][filterX + x * upx];
+                    v *= p.gain;
+                    ((T*)p.y)[outX * p.outStride.x + outY * p.outStride.y + c * p.outStride.z + n * p.outStride.w] = (T)v;
+                }
+            }
+        }
+    }
+}
+
+//------------------------------------------------------------------------
+// CUDA kernel selection.
+
+template <class T> upfirdn2d_kernel_spec choose_upfirdn2d_kernel(const upfirdn2d_kernel_params& p)
+{
+    int s = p.inStride.z, fx = p.filterSize.x, fy = p.filterSize.y;
+
+    upfirdn2d_kernel_spec spec = {(void*)upfirdn2d_kernel_large<T>, -1,-1,1, 4}; // contiguous
+    if (s == 1)           spec = {(void*)upfirdn2d_kernel_large<T>, -1,-1,4, 1}; // channels_last
+
+    if (s != 1 && p.up.x == 1 && p.up.y == 1 && p.down.x == 1 && p.down.y == 1) // contiguous
+    {
+        if (fx <= 7  && fy <= 7 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 7,7,  64,16,1>, 64,16,1, 1};
+        if (fx <= 6  && fy <= 6 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 6,6,  64,16,1>, 64,16,1, 1};
+        if (fx <= 5  && fy <= 5 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 5,5,  64,16,1>, 64,16,1, 1};
+        if (fx <= 4  && fy <= 4 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 4,4,  64,16,1>, 64,16,1, 1};
+        if (fx <= 3  && fy <= 3 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 3,3,  64,16,1>, 64,16,1, 1};
+        if (fx <= 24 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 24,1, 128,8,1>, 128,8,1, 1};
+        if (fx <= 20 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 20,1, 128,8,1>, 128,8,1, 1};
+        if (fx <= 16 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 16,1, 128,8,1>, 128,8,1, 1};
+        if (fx <= 12 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 12,1, 128,8,1>, 128,8,1, 1};
+        if (fx <= 8  && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 8,1,  128,8,1>, 128,8,1, 1};
+        if (fx <= 1  && fy <= 24) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 1,24, 32,32,1>, 32,32,1, 1};
+        if (fx <= 1  && fy <= 20) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 1,20, 32,32,1>, 32,32,1, 1};
+        if (fx <= 1  && fy <= 16) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 1,16, 32,32,1>, 32,32,1, 1};
+        if (fx <= 1  && fy <= 12) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 1,12, 32,32,1>, 32,32,1, 1};
+        if (fx <= 1  && fy <= 8 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 1,8,  32,32,1>, 32,32,1, 1};
+    }
+    if (s == 1 && p.up.x == 1 && p.up.y == 1 && p.down.x == 1 && p.down.y == 1) // channels_last
+    {
+        if (fx <= 7  && fy <= 7 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 7,7,  16,16,8>,  16,16,8,  1};
+        if (fx <= 6  && fy <= 6 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 4,4,  16,16,8>,  16,16,8,  1};
+        if (fx <= 5  && fy <= 5 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 4,4,  16,16,8>,  16,16,8,  1};
+        if (fx <= 4  && fy <= 4 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 4,4,  16,16,8>,  16,16,8,  1};
+        if (fx <= 3  && fy <= 3 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 4,4,  16,16,8>,  16,16,8,  1};
+        if (fx <= 24 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 24,1, 128,1,16>, 128,1,16, 1};
+        if (fx <= 20 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 20,1, 128,1,16>, 128,1,16, 1};
+        if (fx <= 16 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 16,1, 128,1,16>, 128,1,16, 1};
+        if (fx <= 12 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 12,1, 128,1,16>, 128,1,16, 1};
+        if (fx <= 8  && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 8,1,  128,1,16>, 128,1,16, 1};
+        if (fx <= 1  && fy <= 24) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 1,24, 1,128,16>, 1,128,16, 1};
+        if (fx <= 1  && fy <= 20) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 1,20, 1,128,16>, 1,128,16, 1};
+        if (fx <= 1  && fy <= 16) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 1,16, 1,128,16>, 1,128,16, 1};
+        if (fx <= 1  && fy <= 12) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 1,12, 1,128,16>, 1,128,16, 1};
+        if (fx <= 1  && fy <= 8 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,1, 1,8,  1,128,16>, 1,128,16, 1};
+    }
+    if (s != 1 && p.up.x == 2 && p.up.y == 2 && p.down.x == 1 && p.down.y == 1) // contiguous
+    {
+        if (fx <= 8  && fy <= 8 ) spec = {(void*)upfirdn2d_kernel_small<T, 2,2, 1,1, 8,8, 64,16,1>, 64,16,1, 1};
+        if (fx <= 6  && fy <= 6 ) spec = {(void*)upfirdn2d_kernel_small<T, 2,2, 1,1, 6,6, 64,16,1>, 64,16,1, 1};
+        if (fx <= 4  && fy <= 4 ) spec = {(void*)upfirdn2d_kernel_small<T, 2,2, 1,1, 4,4, 64,16,1>, 64,16,1, 1};
+        if (fx <= 2  && fy <= 2 ) spec = {(void*)upfirdn2d_kernel_small<T, 2,2, 1,1, 2,2, 64,16,1>, 64,16,1, 1};
+    }
+    if (s == 1 && p.up.x == 2 && p.up.y == 2 && p.down.x == 1 && p.down.y == 1) // channels_last
+    {
+        if (fx <= 8  && fy <= 8 ) spec = {(void*)upfirdn2d_kernel_small<T, 2,2, 1,1, 8,8, 16,16,8>, 16,16,8, 1};
+        if (fx <= 6  && fy <= 6 ) spec = {(void*)upfirdn2d_kernel_small<T, 2,2, 1,1, 6,6, 16,16,8>, 16,16,8, 1};
+        if (fx <= 4  && fy <= 4 ) spec = {(void*)upfirdn2d_kernel_small<T, 2,2, 1,1, 4,4, 16,16,8>, 16,16,8, 1};
+        if (fx <= 2  && fy <= 2 ) spec = {(void*)upfirdn2d_kernel_small<T, 2,2, 1,1, 2,2, 16,16,8>, 16,16,8, 1};
+    }
+    if (s != 1 && p.up.x == 2 && p.up.y == 1 && p.down.x == 1 && p.down.y == 1) // contiguous
+    {
+        if (fx <= 24 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 2,1, 1,1, 24,1, 128,8,1>, 128,8,1, 1};
+        if (fx <= 20 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 2,1, 1,1, 20,1, 128,8,1>, 128,8,1, 1};
+        if (fx <= 16 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 2,1, 1,1, 16,1, 128,8,1>, 128,8,1, 1};
+        if (fx <= 12 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 2,1, 1,1, 12,1, 128,8,1>, 128,8,1, 1};
+        if (fx <= 8  && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 2,1, 1,1, 8,1,  128,8,1>, 128,8,1, 1};
+    }
+    if (s == 1 && p.up.x == 2 && p.up.y == 1 && p.down.x == 1 && p.down.y == 1) // channels_last
+    {
+        if (fx <= 24 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 2,1, 1,1, 24,1, 128,1,16>, 128,1,16, 1};
+        if (fx <= 20 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 2,1, 1,1, 20,1, 128,1,16>, 128,1,16, 1};
+        if (fx <= 16 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 2,1, 1,1, 16,1, 128,1,16>, 128,1,16, 1};
+        if (fx <= 12 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 2,1, 1,1, 12,1, 128,1,16>, 128,1,16, 1};
+        if (fx <= 8  && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 2,1, 1,1, 8,1,  128,1,16>, 128,1,16, 1};
+    }
+    if (s != 1 && p.up.x == 1 && p.up.y == 2 && p.down.x == 1 && p.down.y == 1) // contiguous
+    {
+        if (fx <= 1  && fy <= 24) spec = {(void*)upfirdn2d_kernel_small<T, 1,2, 1,1, 1,24, 32,32,1>, 32,32,1, 1};
+        if (fx <= 1  && fy <= 20) spec = {(void*)upfirdn2d_kernel_small<T, 1,2, 1,1, 1,20, 32,32,1>, 32,32,1, 1};
+        if (fx <= 1  && fy <= 16) spec = {(void*)upfirdn2d_kernel_small<T, 1,2, 1,1, 1,16, 32,32,1>, 32,32,1, 1};
+        if (fx <= 1  && fy <= 12) spec = {(void*)upfirdn2d_kernel_small<T, 1,2, 1,1, 1,12, 32,32,1>, 32,32,1, 1};
+        if (fx <= 1  && fy <= 8 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,2, 1,1, 1,8,  32,32,1>, 32,32,1, 1};
+    }
+    if (s == 1 && p.up.x == 1 && p.up.y == 2 && p.down.x == 1 && p.down.y == 1) // channels_last
+    {
+        if (fx <= 1  && fy <= 24) spec = {(void*)upfirdn2d_kernel_small<T, 1,2, 1,1, 1,24, 1,128,16>, 1,128,16, 1};
+        if (fx <= 1  && fy <= 20) spec = {(void*)upfirdn2d_kernel_small<T, 1,2, 1,1, 1,20, 1,128,16>, 1,128,16, 1};
+        if (fx <= 1  && fy <= 16) spec = {(void*)upfirdn2d_kernel_small<T, 1,2, 1,1, 1,16, 1,128,16>, 1,128,16, 1};
+        if (fx <= 1  && fy <= 12) spec = {(void*)upfirdn2d_kernel_small<T, 1,2, 1,1, 1,12, 1,128,16>, 1,128,16, 1};
+        if (fx <= 1  && fy <= 8 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,2, 1,1, 1,8,  1,128,16>, 1,128,16, 1};
+    }
+    if (s != 1 && p.up.x == 1 && p.up.y == 1 && p.down.x == 2 && p.down.y == 2) // contiguous
+    {
+        if (fx <= 8  && fy <= 8 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 2,2, 8,8,  32,8,1>, 32,8,1, 1};
+        if (fx <= 6  && fy <= 6 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 2,2, 6,6,  32,8,1>, 32,8,1, 1};
+        if (fx <= 4  && fy <= 4 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 2,2, 4,4,  32,8,1>, 32,8,1, 1};
+        if (fx <= 2  && fy <= 2 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 2,2, 2,2,  32,8,1>, 32,8,1, 1};
+    }
+    if (s == 1 && p.up.x == 1 && p.up.y == 1 && p.down.x == 2 && p.down.y == 2) // channels_last
+    {
+        if (fx <= 8  && fy <= 8 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 2,2, 8,8,  8,8,8>, 8,8,8, 1};
+        if (fx <= 6  && fy <= 6 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 2,2, 6,6,  8,8,8>, 8,8,8, 1};
+        if (fx <= 4  && fy <= 4 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 2,2, 4,4,  8,8,8>, 8,8,8, 1};
+        if (fx <= 2  && fy <= 2 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 2,2, 2,2,  8,8,8>, 8,8,8, 1};
+    }
+    if (s != 1 && p.up.x == 1 && p.up.y == 1 && p.down.x == 2 && p.down.y == 1) // contiguous
+    {
+        if (fx <= 24 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 2,1, 24,1, 64,8,1>, 64,8,1, 1};
+        if (fx <= 20 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 2,1, 20,1, 64,8,1>, 64,8,1, 1};
+        if (fx <= 16 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 2,1, 16,1, 64,8,1>, 64,8,1, 1};
+        if (fx <= 12 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 2,1, 12,1, 64,8,1>, 64,8,1, 1};
+        if (fx <= 8  && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 2,1, 8,1,  64,8,1>, 64,8,1, 1};
+    }
+    if (s == 1 && p.up.x == 1 && p.up.y == 1 && p.down.x == 2 && p.down.y == 1) // channels_last
+    {
+        if (fx <= 24 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 2,1, 24,1, 64,1,8>, 64,1,8, 1};
+        if (fx <= 20 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 2,1, 20,1, 64,1,8>, 64,1,8, 1};
+        if (fx <= 16 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 2,1, 16,1, 64,1,8>, 64,1,8, 1};
+        if (fx <= 12 && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 2,1, 12,1, 64,1,8>, 64,1,8, 1};
+        if (fx <= 8  && fy <= 1 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 2,1, 8,1,  64,1,8>, 64,1,8, 1};
+    }
+    if (s != 1 && p.up.x == 1 && p.up.y == 1 && p.down.x == 1 && p.down.y == 2) // contiguous
+    {
+        if (fx <= 1  && fy <= 24) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,2, 1,24, 32,16,1>, 32,16,1, 1};
+        if (fx <= 1  && fy <= 20) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,2, 1,20, 32,16,1>, 32,16,1, 1};
+        if (fx <= 1  && fy <= 16) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,2, 1,16, 32,16,1>, 32,16,1, 1};
+        if (fx <= 1  && fy <= 12) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,2, 1,12, 32,16,1>, 32,16,1, 1};
+        if (fx <= 1  && fy <= 8 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,2, 1,8,  32,16,1>, 32,16,1, 1};
+    }
+    if (s == 1 && p.up.x == 1 && p.up.y == 1 && p.down.x == 1 && p.down.y == 2) // channels_last
+    {
+        if (fx <= 1  && fy <= 24) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,2, 1,24, 1,64,8>, 1,64,8, 1};
+        if (fx <= 1  && fy <= 20) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,2, 1,20, 1,64,8>, 1,64,8, 1};
+        if (fx <= 1  && fy <= 16) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,2, 1,16, 1,64,8>, 1,64,8, 1};
+        if (fx <= 1  && fy <= 12) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,2, 1,12, 1,64,8>, 1,64,8, 1};
+        if (fx <= 1  && fy <= 8 ) spec = {(void*)upfirdn2d_kernel_small<T, 1,1, 1,2, 1,8,  1,64,8>, 1,64,8, 1};
+    }
+    return spec;
+}
+
+//------------------------------------------------------------------------
+// Template specializations.
+
+template upfirdn2d_kernel_spec choose_upfirdn2d_kernel<double>   (const upfirdn2d_kernel_params& p);
+template upfirdn2d_kernel_spec choose_upfirdn2d_kernel<float>    (const upfirdn2d_kernel_params& p);
+template upfirdn2d_kernel_spec choose_upfirdn2d_kernel<c10::Half>(const upfirdn2d_kernel_params& p);
+
+//------------------------------------------------------------------------

third_party/stylegan2_official_ops/upfirdn2d.h

@@ -0,0 +1,59 @@
+// 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
+// and any modifications thereto.  Any use, reproduction, disclosure or
+// distribution of this software and related documentation without an express
+// license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+#include <cuda_runtime.h>
+
+//------------------------------------------------------------------------
+// CUDA kernel parameters.
+
+struct upfirdn2d_kernel_params
+{
+    const void*     x;
+    const float*    f;
+    void*           y;
+
+    int2            up;
+    int2            down;
+    int2            pad0;
+    int             flip;
+    float           gain;
+
+    int4            inSize;         // [width, height, channel, batch]
+    int4            inStride;
+    int2            filterSize;     // [width, height]
+    int2            filterStride;
+    int4            outSize;        // [width, height, channel, batch]
+    int4            outStride;
+    int             sizeMinor;
+    int             sizeMajor;
+
+    int             loopMinor;
+    int             loopMajor;
+    int             loopX;
+    int             launchMinor;
+    int             launchMajor;
+};
+
+//------------------------------------------------------------------------
+// CUDA kernel specialization.
+
+struct upfirdn2d_kernel_spec
+{
+    void*   kernel;
+    int     tileOutW;
+    int     tileOutH;
+    int     loopMinor;
+    int     loopX;
+};
+
+//------------------------------------------------------------------------
+// CUDA kernel selection.
+
+template <class T> upfirdn2d_kernel_spec choose_upfirdn2d_kernel(const upfirdn2d_kernel_params& p);
+
+//------------------------------------------------------------------------

third_party/stylegan2_official_ops/upfirdn2d.py

@@ -0,0 +1,401 @@
+# python3.7
+
+# 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
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+"""Custom operators for efficient resampling of 2D images.
+
+`upfirdn` means executing upsampling, FIR filtering, downsampling in sequence.
+
+Please refer to https://github.com/NVlabs/stylegan2-ada-pytorch
+"""
+
+# pylint: disable=line-too-long
+# pylint: disable=missing-class-docstring
+# pylint: disable=global-variable-not-assigned
+# pylint: disable=bare-except
+
+import os
+import warnings
+import traceback
+import numpy as np
+import torch
+
+from . import custom_ops
+from . import misc
+from . import conv2d_gradfix
+
+#----------------------------------------------------------------------------
+
+_inited = False
+_plugin = None
+
+def _init():
+    global _inited, _plugin
+    if not _inited:
+        sources = ['upfirdn2d.cpp', 'upfirdn2d.cu']
+        sources = [os.path.join(os.path.dirname(__file__), s) for s in sources]
+        try:
+            _plugin = custom_ops.get_plugin('upfirdn2d_plugin', sources=sources, extra_cuda_cflags=['--use_fast_math'])
+        except:
+            warnings.warn('Failed to build CUDA kernels for upfirdn2d. Falling back to slow reference implementation. Details:\n\n' + traceback.format_exc())
+    return _plugin is not None
+
+def _parse_scaling(scaling):
+    if isinstance(scaling, int):
+        scaling = [scaling, scaling]
+    assert isinstance(scaling, (list, tuple))
+    assert all(isinstance(x, int) for x in scaling)
+    sx, sy = scaling
+    assert sx >= 1 and sy >= 1
+    return sx, sy
+
+def _parse_padding(padding):
+    if isinstance(padding, int):
+        padding = [padding, padding]
+    assert isinstance(padding, (list, tuple))
+    assert all(isinstance(x, int) for x in padding)
+    if len(padding) == 2:
+        padx, pady = padding
+        padding = [padx, padx, pady, pady]
+    padx0, padx1, pady0, pady1 = padding
+    return padx0, padx1, pady0, pady1
+
+def _get_filter_size(f):
+    if f is None:
+        return 1, 1
+    assert isinstance(f, torch.Tensor) and f.ndim in [1, 2]
+    fw = f.shape[-1]
+    fh = f.shape[0]
+    with misc.suppress_tracer_warnings():
+        fw = int(fw)
+        fh = int(fh)
+    misc.assert_shape(f, [fh, fw][:f.ndim])
+    assert fw >= 1 and fh >= 1
+    return fw, fh
+
+#----------------------------------------------------------------------------
+
+def setup_filter(f, device=torch.device('cpu'), normalize=True, flip_filter=False, gain=1, separable=None):
+    r"""Convenience function to setup 2D FIR filter for `upfirdn2d()`.
+
+    Args:
+        f:           Torch tensor, numpy array, or python list of the shape
+                     `[filter_height, filter_width]` (non-separable),
+                     `[filter_taps]` (separable),
+                     `[]` (impulse), or
+                     `None` (identity).
+        device:      Result device (default: cpu).
+        normalize:   Normalize the filter so that it retains the magnitude
+                     for constant input signal (DC)? (default: True).
+        flip_filter: Flip the filter? (default: False).
+        gain:        Overall scaling factor for signal magnitude (default: 1).
+        separable:   Return a separable filter? (default: select automatically).
+
+    Returns:
+        Float32 tensor of the shape
+        `[filter_height, filter_width]` (non-separable) or
+        `[filter_taps]` (separable).
+    """
+    # Validate.
+    if f is None:
+        f = 1
+    f = torch.as_tensor(f, dtype=torch.float32)
+    assert f.ndim in [0, 1, 2]
+    assert f.numel() > 0
+    if f.ndim == 0:
+        f = f[np.newaxis]
+
+    # Separable?
+    if separable is None:
+        separable = (f.ndim == 1 and f.numel() >= 8)
+    if f.ndim == 1 and not separable:
+        f = f.ger(f)
+    assert f.ndim == (1 if separable else 2)
+
+    # Apply normalize, flip, gain, and device.
+    if normalize:
+        f /= f.sum()
+    if flip_filter:
+        f = f.flip(list(range(f.ndim)))
+    f = f * (gain ** (f.ndim / 2))
+    f = f.to(device=device)
+    return f
+
+#----------------------------------------------------------------------------
+
+def upfirdn2d(x, f, up=1, down=1, padding=0, flip_filter=False, gain=1, impl='cuda'):
+    r"""Pad, upsample, filter, and downsample a batch of 2D images.
+
+    Performs the following sequence of operations for each channel:
+
+    1. Upsample the image by inserting N-1 zeros after each pixel (`up`).
+
+    2. Pad the image with the specified number of zeros on each side (`padding`).
+       Negative padding corresponds to cropping the image.
+
+    3. Convolve the image with the specified 2D FIR filter (`f`), shrinking it
+       so that the footprint of all output pixels lies within the input image.
+
+    4. Downsample the image by keeping every Nth pixel (`down`).
+
+    This sequence of operations bears close resemblance to scipy.signal.upfirdn().
+    The fused op is considerably more efficient than performing the same calculation
+    using standard PyTorch ops. It supports gradients of arbitrary order.
+
+    Args:
+        x:           Float32/float64/float16 input tensor of the shape
+                     `[batch_size, num_channels, in_height, in_width]`.
+        f:           Float32 FIR filter of the shape
+                     `[filter_height, filter_width]` (non-separable),
+                     `[filter_taps]` (separable), or
+                     `None` (identity).
+        up:          Integer upsampling factor. Can be a single int or a list/tuple
+                     `[x, y]` (default: 1).
+        down:        Integer downsampling factor. Can be a single int or a list/tuple
+                     `[x, y]` (default: 1).
+        padding:     Padding with respect to the upsampled image. Can be a single number
+                     or a list/tuple `[x, y]` or `[x_before, x_after, y_before, y_after]`
+                     (default: 0).
+        flip_filter: False = convolution, True = correlation (default: False).
+        gain:        Overall scaling factor for signal magnitude (default: 1).
+        impl:        Implementation to use. Can be `'ref'` or `'cuda'` (default: `'cuda'`).
+
+    Returns:
+        Tensor of the shape `[batch_size, num_channels, out_height, out_width]`.
+    """
+    assert isinstance(x, torch.Tensor)
+    assert impl in ['ref', 'cuda']
+    if impl == 'cuda' and x.device.type == 'cuda' and _init():
+        return _upfirdn2d_cuda(up=up, down=down, padding=padding, flip_filter=flip_filter, gain=gain).apply(x, f)
+    return _upfirdn2d_ref(x, f, up=up, down=down, padding=padding, flip_filter=flip_filter, gain=gain)
+
+#----------------------------------------------------------------------------
+
+@misc.profiled_function
+def _upfirdn2d_ref(x, f, up=1, down=1, padding=0, flip_filter=False, gain=1):
+    """Slow reference implementation of `upfirdn2d()` using standard PyTorch ops.
+    """
+    # Validate arguments.
+    assert isinstance(x, torch.Tensor) and x.ndim == 4
+    if f is None:
+        f = torch.ones([1, 1], dtype=torch.float32, device=x.device)
+    assert isinstance(f, torch.Tensor) and f.ndim in [1, 2]
+    assert f.dtype == torch.float32 and not f.requires_grad
+    batch_size, num_channels, in_height, in_width = x.shape
+    upx, upy = _parse_scaling(up)
+    downx, downy = _parse_scaling(down)
+    padx0, padx1, pady0, pady1 = _parse_padding(padding)
+
+    # Upsample by inserting zeros.
+    x = x.reshape([batch_size, num_channels, in_height, 1, in_width, 1])
+    x = torch.nn.functional.pad(x, [0, upx - 1, 0, 0, 0, upy - 1])
+    x = x.reshape([batch_size, num_channels, in_height * upy, in_width * upx])
+
+    # Pad or crop.
+    x = torch.nn.functional.pad(x, [max(padx0, 0), max(padx1, 0), max(pady0, 0), max(pady1, 0)])
+    x = x[:, :, max(-pady0, 0) : x.shape[2] - max(-pady1, 0), max(-padx0, 0) : x.shape[3] - max(-padx1, 0)]
+
+    # Setup filter.
+    f = f * (gain ** (f.ndim / 2))
+    f = f.to(x.dtype)
+    if not flip_filter:
+        f = f.flip(list(range(f.ndim)))
+
+    # Convolve with the filter.
+    f = f[np.newaxis, np.newaxis].repeat([num_channels, 1] + [1] * f.ndim)
+    if f.ndim == 4:
+        x = conv2d_gradfix.conv2d(input=x, weight=f, groups=num_channels)
+    else:
+        x = conv2d_gradfix.conv2d(input=x, weight=f.unsqueeze(2), groups=num_channels)
+        x = conv2d_gradfix.conv2d(input=x, weight=f.unsqueeze(3), groups=num_channels)
+
+    # Downsample by throwing away pixels.
+    x = x[:, :, ::downy, ::downx]
+    return x
+
+#----------------------------------------------------------------------------
+
+_upfirdn2d_cuda_cache = dict()
+
+def _upfirdn2d_cuda(up=1, down=1, padding=0, flip_filter=False, gain=1):
+    """Fast CUDA implementation of `upfirdn2d()` using custom ops.
+    """
+    # Parse arguments.
+    upx, upy = _parse_scaling(up)
+    downx, downy = _parse_scaling(down)
+    padx0, padx1, pady0, pady1 = _parse_padding(padding)
+
+    # Lookup from cache.
+    key = (upx, upy, downx, downy, padx0, padx1, pady0, pady1, flip_filter, gain)
+    if key in _upfirdn2d_cuda_cache:
+        return _upfirdn2d_cuda_cache[key]
+
+    # Forward op.
+    class Upfirdn2dCuda(torch.autograd.Function):
+        @staticmethod
+        def forward(ctx, x, f): # pylint: disable=arguments-differ
+            assert isinstance(x, torch.Tensor) and x.ndim == 4
+            if f is None:
+                f = torch.ones([1, 1], dtype=torch.float32, device=x.device)
+            assert isinstance(f, torch.Tensor) and f.ndim in [1, 2]
+            y = x
+            if f.ndim == 2:
+                y = _plugin.upfirdn2d(y, f, upx, upy, downx, downy, padx0, padx1, pady0, pady1, flip_filter, gain)
+            else:
+                y = _plugin.upfirdn2d(y, f.unsqueeze(0), upx, 1, downx, 1, padx0, padx1, 0, 0, flip_filter, np.sqrt(gain))
+                y = _plugin.upfirdn2d(y, f.unsqueeze(1), 1, upy, 1, downy, 0, 0, pady0, pady1, flip_filter, np.sqrt(gain))
+            ctx.save_for_backward(f)
+            ctx.x_shape = x.shape
+            return y
+
+        @staticmethod
+        def backward(ctx, dy): # pylint: disable=arguments-differ
+            f, = ctx.saved_tensors
+            _, _, ih, iw = ctx.x_shape
+            _, _, oh, ow = dy.shape
+            fw, fh = _get_filter_size(f)
+            p = [
+                fw - padx0 - 1,
+                iw * upx - ow * downx + padx0 - upx + 1,
+                fh - pady0 - 1,
+                ih * upy - oh * downy + pady0 - upy + 1,
+            ]
+            dx = None
+            df = None
+
+            if ctx.needs_input_grad[0]:
+                dx = _upfirdn2d_cuda(up=down, down=up, padding=p, flip_filter=(not flip_filter), gain=gain).apply(dy, f)
+
+            assert not ctx.needs_input_grad[1]
+            return dx, df
+
+    # Add to cache.
+    _upfirdn2d_cuda_cache[key] = Upfirdn2dCuda
+    return Upfirdn2dCuda
+
+#----------------------------------------------------------------------------
+
+def filter2d(x, f, padding=0, flip_filter=False, gain=1, impl='cuda'):
+    r"""Filter a batch of 2D images using the given 2D FIR filter.
+
+    By default, the result is padded so that its shape matches the input.
+    User-specified padding is applied on top of that, with negative values
+    indicating cropping. Pixels outside the image are assumed to be zero.
+
+    Args:
+        x:           Float32/float64/float16 input tensor of the shape
+                     `[batch_size, num_channels, in_height, in_width]`.
+        f:           Float32 FIR filter of the shape
+                     `[filter_height, filter_width]` (non-separable),
+                     `[filter_taps]` (separable), or
+                     `None` (identity).
+        padding:     Padding with respect to the output. Can be a single number or a
+                     list/tuple `[x, y]` or `[x_before, x_after, y_before, y_after]`
+                     (default: 0).
+        flip_filter: False = convolution, True = correlation (default: False).
+        gain:        Overall scaling factor for signal magnitude (default: 1).
+        impl:        Implementation to use. Can be `'ref'` or `'cuda'` (default: `'cuda'`).
+
+    Returns:
+        Tensor of the shape `[batch_size, num_channels, out_height, out_width]`.
+    """
+    padx0, padx1, pady0, pady1 = _parse_padding(padding)
+    fw, fh = _get_filter_size(f)
+    p = [
+        padx0 + fw // 2,
+        padx1 + (fw - 1) // 2,
+        pady0 + fh // 2,
+        pady1 + (fh - 1) // 2,
+    ]
+    return upfirdn2d(x, f, padding=p, flip_filter=flip_filter, gain=gain, impl=impl)
+
+#----------------------------------------------------------------------------
+
+def upsample2d(x, f, up=2, padding=0, flip_filter=False, gain=1, impl='cuda'):
+    r"""Upsample a batch of 2D images using the given 2D FIR filter.
+
+    By default, the result is padded so that its shape is a multiple of the input.
+    User-specified padding is applied on top of that, with negative values
+    indicating cropping. Pixels outside the image are assumed to be zero.
+
+    Args:
+        x:           Float32/float64/float16 input tensor of the shape
+                     `[batch_size, num_channels, in_height, in_width]`.
+        f:           Float32 FIR filter of the shape
+                     `[filter_height, filter_width]` (non-separable),
+                     `[filter_taps]` (separable), or
+                     `None` (identity).
+        up:          Integer upsampling factor. Can be a single int or a list/tuple
+                     `[x, y]` (default: 1).
+        padding:     Padding with respect to the output. Can be a single number or a
+                     list/tuple `[x, y]` or `[x_before, x_after, y_before, y_after]`
+                     (default: 0).
+        flip_filter: False = convolution, True = correlation (default: False).
+        gain:        Overall scaling factor for signal magnitude (default: 1).
+        impl:        Implementation to use. Can be `'ref'` or `'cuda'` (default: `'cuda'`).
+
+    Returns:
+        Tensor of the shape `[batch_size, num_channels, out_height, out_width]`.
+    """
+    upx, upy = _parse_scaling(up)
+    padx0, padx1, pady0, pady1 = _parse_padding(padding)
+    fw, fh = _get_filter_size(f)
+    p = [
+        padx0 + (fw + upx - 1) // 2,
+        padx1 + (fw - upx) // 2,
+        pady0 + (fh + upy - 1) // 2,
+        pady1 + (fh - upy) // 2,
+    ]
+    return upfirdn2d(x, f, up=up, padding=p, flip_filter=flip_filter, gain=gain*upx*upy, impl=impl)
+
+#----------------------------------------------------------------------------
+
+def downsample2d(x, f, down=2, padding=0, flip_filter=False, gain=1, impl='cuda'):
+    r"""Downsample a batch of 2D images using the given 2D FIR filter.
+
+    By default, the result is padded so that its shape is a fraction of the input.
+    User-specified padding is applied on top of that, with negative values
+    indicating cropping. Pixels outside the image are assumed to be zero.
+
+    Args:
+        x:           Float32/float64/float16 input tensor of the shape
+                     `[batch_size, num_channels, in_height, in_width]`.
+        f:           Float32 FIR filter of the shape
+                     `[filter_height, filter_width]` (non-separable),
+                     `[filter_taps]` (separable), or
+                     `None` (identity).
+        down:        Integer downsampling factor. Can be a single int or a list/tuple
+                     `[x, y]` (default: 1).
+        padding:     Padding with respect to the input. Can be a single number or a
+                     list/tuple `[x, y]` or `[x_before, x_after, y_before, y_after]`
+                     (default: 0).
+        flip_filter: False = convolution, True = correlation (default: False).
+        gain:        Overall scaling factor for signal magnitude (default: 1).
+        impl:        Implementation to use. Can be `'ref'` or `'cuda'` (default: `'cuda'`).
+
+    Returns:
+        Tensor of the shape `[batch_size, num_channels, out_height, out_width]`.
+    """
+    downx, downy = _parse_scaling(down)
+    padx0, padx1, pady0, pady1 = _parse_padding(padding)
+    fw, fh = _get_filter_size(f)
+    p = [
+        padx0 + (fw - downx + 1) // 2,
+        padx1 + (fw - downx) // 2,
+        pady0 + (fh - downy + 1) // 2,
+        pady1 + (fh - downy) // 2,
+    ]
+    return upfirdn2d(x, f, down=down, padding=p, flip_filter=flip_filter, gain=gain, impl=impl)
+
+#----------------------------------------------------------------------------
+
+# pylint: enable=line-too-long
+# pylint: enable=missing-class-docstring
+# pylint: enable=global-variable-not-assigned
+# pylint: enable=bare-except

third_party/stylegan3_official_ops/README.md

@@ -0,0 +1,30 @@
+# Operators for StyleGAN2
+
+All files in this directory are borrowed from repository [stylegan3](https://github.com/NVlabs/stylegan3). Basically, these files implement customized operators, which are faster than the native operators from PyTorch, especially for second-derivative computation, including
+
+- `bias_act.bias_act()`: Fuse adding bias and then performing activation as one operator.
+- `upfirdn2d.setup_filter()`: Set up the kernel used for filtering.
+- `upfirdn2d.filter2d()`: Filtering a 2D feature map with given kernel.
+- `upfirdn2d.upsample2d()`: Upsampling a 2D feature map.
+- `upfirdn2d.downsample2d()`: Downsampling a 2D feature map.
+- `upfirdn2d.upfirdn2d()`: Upsampling, filtering, and then downsampling a 2D feature map.
+- `filtered_lrelu.filtered_lrelu()`: Leaky ReLU layer, wrapped with upsampling and downsampling for anti-aliasing.
+- `conv2d_gradfix.conv2d()`: Convolutional layer, supporting arbitrarily high order gradients and fixing gradient when computing penalty.
+- `conv2d_gradfix.conv_transpose2d()`: Transposed convolutional layer, supporting arbitrarily high order gradients and fixing gradient when computing penalty.
+- `conv2d_resample.conv2d_resample()`: Wraps `upfirdn2d()` and `conv2d()` (or `conv_transpose2d()`). This is not used in our network implementation (*i.e.*, `models/stylegan2_generator.py` and `models/stylegan2_discriminator.py`)
+
+We make following slight modifications beyond disabling some lint warnings:
+
+- Line 24 of file `misc.py`: Use `EasyDict` from module `easydict` to replace that from `dnnlib` from [stylegan3](https://github.com/NVlabs/stylegan3).
+- Line 36 of file `custom_ops.py`: Disable log message when setting up customized operators.
+- Line 54/109 of file `custom_ops.py`: Add necessary CUDA compiler path. (***NOTE**: If your cuda binary does not locate at `/usr/local/cuda/bin`, please specify in function `_find_compiler_bindir_posix()`.*)
+- Line 21 of file `bias_act.py`: Use `EasyDict` from module `easydict` to replace that from `dnnlib` from [stylegan3](https://github.com/NVlabs/stylegan3).
+- Line 162-165 of file `filtered_lrelu.py`: Change some implementations in `_filtered_lrelu_ref()` to `ref`.
+- Line 31 of file `grid_sample_gradfix.py`: Enable customized grid sampling operator by default.
+- Line 35 of file `grid_sample_gradfix.py`:  Use `impl` to disable customized grid sample operator.
+- Line 34 of file `conv2d_gradfix.py`: Enable customized convolution operators by default.
+- Line 48/53 of file `conv2d_gradfix.py`: Use `impl` to disable customized convolution operators.
+- Line 36/53 of file `conv2d_resample.py`: Use `impl` to disable customized convolution operators.
+- Line 23 of file `fma.py`: Use `impl` to disable customized add-multiply operator.
+
+Please use `ref` or `cuda` to choose which implementation to use. `ref` refers to native PyTorch operators while `cuda` refers to the customized operators from the official repository. `cuda` is used by default.