calc_fid.py

# Copyright (C) 2021 NVIDIA Corporation.  All rights reserved.
# Licensed under The MIT License (MIT)
# Permission is hereby granted, free of charge, to any person obtaining a copy of
# this software and associated documentation files (the "Software"), to deal in
# the Software without restriction, including without limitation the rights to
# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
# the Software, and to permit persons to whom the Software is furnished to do so,
# subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

import os
import argparse
import torch
from models import make_model

import functools
from utils.inception_utils import sample_gema, prepare_inception_metrics

if __name__ == "__main__":
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    parser = argparse.ArgumentParser(
        description="Calculate FID score for generators",
    )
    parser.add_argument(
        "--ckpt",
        type=str,
        required=True,
        help="path to the checkpoint file",
    )
    parser.add_argument(
        "--inception",
        type=str,
        required=True,
        help="pre-calculated inception file",
    )
    parser.add_argument(
        "--batch", default=8, type=int, help="batch size for inception networks"
    )
    parser.add_argument(
        "--n_sample",
        type=int,
        default=50000,
        help="number of samples used for embedding calculation",
    )
    args = parser.parse_args()

    print("Loading model...")
    ckpt = torch.load(args.ckpt)
    g_args = ckpt['args']
    model = make_model(g_args).to(device).eval()
    model.load_state_dict(ckpt['g_ema'])
    mean_latent = model.style(torch.randn(50000,512, device=device)).mean(0)

    get_inception_metrics = prepare_inception_metrics(args.inception, False)
    sample_fn = functools.partial(sample_gema, g_ema=model, device=device, 
                    truncation=1.0, mean_latent=None, batch_size=args.batch)

    print("==================Start calculating FID==================")
    IS_mean, IS_std, FID = get_inception_metrics(sample_fn, num_inception_images=args.n_sample, use_torch=False)
    print("FID: {0:.4f}, IS_mean: {1:.4f}, IS_std: {2:.4f}".format(FID, IS_mean, IS_std))