calc_metrics.py

# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES.  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.

"""Calculate quality metrics for previous training run or pretrained network pickle."""

import copy
import json
import os
import tempfile

import click
import dnnlib
import legacy
import torch
from metrics import metric_main, metric_utils
from torch_utils import custom_ops, misc, training_stats
from torch_utils.ops import conv2d_gradfix


# ----------------------------------------------------------------------------


def subprocess_fn(rank, args, temp_dir):
    dnnlib.util.Logger(should_flush=True)

    # Init torch.distributed.
    if args.num_gpus > 1:
        init_file = os.path.abspath(os.path.join(temp_dir, ".torch_distributed_init"))
        if os.name == "nt":
            init_method = "file:///" + init_file.replace("\\", "/")
            torch.distributed.init_process_group(
                backend="gloo",
                init_method=init_method,
                rank=rank,
                world_size=args.num_gpus,
            )
        else:
            init_method = f"file://{init_file}"
            torch.distributed.init_process_group(
                backend="nccl",
                init_method=init_method,
                rank=rank,
                world_size=args.num_gpus,
            )

    # Init torch_utils.
    sync_device = torch.device("cuda", rank) if args.num_gpus > 1 else None
    training_stats.init_multiprocessing(rank=rank, sync_device=sync_device)
    if rank != 0 or not args.verbose:
        custom_ops.verbosity = "none"

    # Configure torch.
    device = torch.device("cuda", rank)
    torch.backends.cuda.matmul.allow_tf32 = False
    torch.backends.cudnn.allow_tf32 = False
    conv2d_gradfix.enabled = True

    # Print network summary.
    G = copy.deepcopy(args.G).eval().requires_grad_(False).to(device)
    if rank == 0 and args.verbose:
        z = torch.empty([1, G.z_dim], device=device)
        c = torch.empty([1, G.c_dim], device=device)
        misc.print_module_summary(G, [z, c])

    # Calculate each metric.
    for metric in args.metrics:
        if rank == 0 and args.verbose:
            print(f"Calculating {metric}...")
        progress = metric_utils.ProgressMonitor(verbose=args.verbose)
        result_dict = metric_main.calc_metric(
            metric=metric,
            G=G,
            dataset_kwargs=args.dataset_kwargs,
            num_gpus=args.num_gpus,
            rank=rank,
            device=device,
            progress=progress,
        )
        if rank == 0:
            metric_main.report_metric(result_dict, run_dir=args.run_dir, snapshot_pkl=args.network_pkl)
        if rank == 0 and args.verbose:
            print()

    # Done.
    if rank == 0 and args.verbose:
        print("Exiting...")


# ----------------------------------------------------------------------------


def parse_comma_separated_list(s):
    if isinstance(s, list):
        return s
    if s is None or s.lower() == "none" or s == "":
        return []
    return s.split(",")


# ----------------------------------------------------------------------------


def make_coords(resolution: float, scale: float):
    coords = torch.linspace(0, 1, int(resolution * scale))
    coords = coords.reshape(1, -1, 1, 1)
    coords = coords.repeat(1, 1, 2, 1)
    return coords


# ----------------------------------------------------------------------------


@click.command()
@click.pass_context
@click.option(
    "network_pkl",
    "--network",
    help="Network pickle filename or URL",
    metavar="PATH",
    required=True,
)
@click.option(
    "--scale",
    help="Scale of generated images",
    type=float,
    default=1,
    show_default=True,
)
@click.option(
    "--metrics",
    help="Quality metrics",
    metavar="[NAME|A,B,C|none]",
    type=parse_comma_separated_list,
    default="fid50k_full",
    show_default=True,
)
@click.option(
    "--data",
    help="Dataset to evaluate against  [default: look up]",
    metavar="[ZIP|DIR]",
)
@click.option(
    "--mirror",
    help="Enable dataset x-flips  [default: look up]",
    type=bool,
    metavar="BOOL",
)
@click.option(
    "--gpus",
    help="Number of GPUs to use",
    type=int,
    default=1,
    metavar="INT",
    show_default=True,
)
@click.option(
    "--verbose",
    help="Print optional information",
    type=bool,
    default=True,
    metavar="BOOL",
    show_default=True,
)
def calc_metrics(ctx, network_pkl, scale, metrics, data, mirror, gpus, verbose):
    """Calculate quality metrics for previous training run or pretrained network pickle.

    Examples:

    \b
    # Previous training run: look up options automatically, save result to JSONL file.
    python calc_metrics.py --metrics=eqt50k_int,eqr50k \\
        --network=~/training-runs/00000-stylegan3-r-mydataset/network-snapshot-000000.pkl

    \b
    # Pre-trained network pickle: specify dataset explicitly, print result to stdout.
    python calc_metrics.py --metrics=fid50k_full --data=~/datasets/ffhq-1024x1024.zip --mirror=1 \\
        --network=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-t-ffhq-1024x1024.pkl

    \b
    Recommended metrics:
      fid50k_full  Frechet inception distance against the full dataset.
      kid50k_full  Kernel inception distance against the full dataset.
      pr50k3_full  Precision and recall againt the full dataset.
      ppl2_wend    Perceptual path length in W, endpoints, full image.
      eqt50k_int   Equivariance w.r.t. integer translation (EQ-T).
      eqt50k_frac  Equivariance w.r.t. fractional translation (EQ-T_frac).
      eqr50k       Equivariance w.r.t. rotation (EQ-R).

    \b
    Legacy metrics:
      fid50k       Frechet inception distance against 50k real images.
      kid50k       Kernel inception distance against 50k real images.
      pr50k3       Precision and recall against 50k real images.
      is50k        Inception score for CIFAR-10.
    """
    dnnlib.util.Logger(should_flush=True)

    # Validate arguments.
    args = dnnlib.EasyDict(metrics=metrics, num_gpus=gpus, network_pkl=network_pkl, verbose=verbose)
    if not all(metric_main.is_valid_metric(metric) for metric in args.metrics):
        ctx.fail("\n".join(["--metrics can only contain the following values:"] + metric_main.list_valid_metrics()))
    if not args.num_gpus >= 1:
        ctx.fail("--gpus must be at least 1")

    # Load network.
    if not dnnlib.util.is_url(network_pkl, allow_file_urls=True) and not os.path.isfile(network_pkl):
        ctx.fail("--network must point to a file or URL")
    if args.verbose:
        print(f'Loading network from "{network_pkl}"...')
    with dnnlib.util.open_url(network_pkl, verbose=args.verbose) as f:
        network_dict = legacy.load_network_pkl(f)
        args.G = network_dict["G_ema"]  # subclass of torch.nn.Module
        
    # Construct an input coordinates and pass to the creps generator.
    if hasattr(args.G.synthesis.b4, "input"):
        args.G.synthesis.b4.input.coords = make_coords(args.G.img_resolution, scale)

    # Initialize dataset options.
    if data is not None:
        args.dataset_kwargs = dnnlib.EasyDict(class_name="training.dataset.ImageFolderDataset", path=data)
    elif network_dict["training_set_kwargs"] is not None:
        args.dataset_kwargs = dnnlib.EasyDict(network_dict["training_set_kwargs"])
    else:
        ctx.fail("Could not look up dataset options; please specify --data")

    # Finalize dataset options.
    args.dataset_kwargs.resolution = int(args.G.img_resolution * scale)
    args.dataset_kwargs.use_labels = args.G.c_dim != 0
    if mirror is not None:
        args.dataset_kwargs.xflip = mirror

    # Print dataset options.
    if args.verbose:
        print("Dataset options:")
        print(json.dumps(args.dataset_kwargs, indent=2))

    # Locate run dir.
    args.run_dir = None
    if os.path.isfile(network_pkl):
        pkl_dir = os.path.dirname(network_pkl)
        if os.path.isfile(os.path.join(pkl_dir, "training_options.json")):
            args.run_dir = pkl_dir

    # Launch processes.
    if args.verbose:
        print("Launching processes...")
    torch.multiprocessing.set_start_method("spawn")
    with tempfile.TemporaryDirectory() as temp_dir:
        if args.num_gpus == 1:
            subprocess_fn(rank=0, args=args, temp_dir=temp_dir)
        else:
            torch.multiprocessing.spawn(fn=subprocess_fn, args=(args, temp_dir), nprocs=args.num_gpus)


# ----------------------------------------------------------------------------

if __name__ == "__main__":
    calc_metrics()  # pylint: disable=no-value-for-parameter

# ----------------------------------------------------------------------------