update repo

.gitattributes

@@ -34,3 +34,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 *.pth* filter=lfs diff=lfs merge=lfs -text
  filter=lfs diff=lfs merge=lfs -text
+*.pkl* filter=lfs diff=lfs merge=lfs -text
+ filter=lfs diff=lfs merge=lfs -text

id_loss.py

@@ -15,7 +15,7 @@ class IDLoss(nn.Module):
         super(IDLoss, self).__init__()
         print('Loading ResNet ArcFace')
         self.facenet = Backbone(input_size=112, num_layers=50, drop_ratio=0.6, mode='ir_se')
-        self.facenet.load_state_dict(torch.load("model_ir_se50.pth", map_location=device))
+        self.facenet.load_state_dict(torch.load("./pretrained/model_ir_se50.pth", map_location=device))
         self.pool = torch.nn.AdaptiveAvgPool2d((256, 256))
         self.face_pool = torch.nn.AdaptiveAvgPool2d((112, 112))
         self.facenet.eval()

metrics/__init__.py

@@ -1,9 +0,0 @@
-# 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.
-
-# empty

metrics/frechet_inception_distance.py

@@ -1,41 +0,0 @@
-# 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.
-
-"""Frechet Inception Distance (FID) from the paper
-"GANs trained by a two time-scale update rule converge to a local Nash
-equilibrium". Matches the original implementation by Heusel et al. at
-https://github.com/bioinf-jku/TTUR/blob/master/fid.py"""
-
-import numpy as np
-import scipy.linalg
-from . import metric_utils
-
-#----------------------------------------------------------------------------
-
-def compute_fid(opts, max_real, num_gen):
-    # Direct TorchScript translation of http://download.tensorflow.org/models/image/imagenet/inception-2015-12-05.tgz
-    detector_url = 'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metrics/inception-2015-12-05.pt'
-    detector_kwargs = dict(return_features=True) # Return raw features before the softmax layer.
-
-    mu_real, sigma_real = metric_utils.compute_feature_stats_for_dataset(
-        opts=opts, detector_url=detector_url, detector_kwargs=detector_kwargs,
-        rel_lo=0, rel_hi=0, capture_mean_cov=True, max_items=max_real).get_mean_cov()
-
-    mu_gen, sigma_gen = metric_utils.compute_feature_stats_for_generator(
-        opts=opts, detector_url=detector_url, detector_kwargs=detector_kwargs,
-        rel_lo=0, rel_hi=1, capture_mean_cov=True, max_items=num_gen).get_mean_cov()
-
-    if opts.rank != 0:
-        return float('nan')
-
-    m = np.square(mu_gen - mu_real).sum()
-    s, _ = scipy.linalg.sqrtm(np.dot(sigma_gen, sigma_real), disp=False) # pylint: disable=no-member
-    fid = np.real(m + np.trace(sigma_gen + sigma_real - s * 2))
-    return float(fid)
-
-#----------------------------------------------------------------------------

metrics/inception_score.py

@@ -1,38 +0,0 @@
-# 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.
-
-"""Inception Score (IS) from the paper "Improved techniques for training
-GANs". Matches the original implementation by Salimans et al. at
-https://github.com/openai/improved-gan/blob/master/inception_score/model.py"""
-
-import numpy as np
-from . import metric_utils
-
-#----------------------------------------------------------------------------
-
-def compute_is(opts, num_gen, num_splits):
-    # Direct TorchScript translation of http://download.tensorflow.org/models/image/imagenet/inception-2015-12-05.tgz
-    detector_url = 'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metrics/inception-2015-12-05.pt'
-    detector_kwargs = dict(no_output_bias=True) # Match the original implementation by not applying bias in the softmax layer.
-
-    gen_probs = metric_utils.compute_feature_stats_for_generator(
-        opts=opts, detector_url=detector_url, detector_kwargs=detector_kwargs,
-        capture_all=True, max_items=num_gen).get_all()
-
-    if opts.rank != 0:
-        return float('nan'), float('nan')
-
-    scores = []
-    for i in range(num_splits):
-        part = gen_probs[i * num_gen // num_splits : (i + 1) * num_gen // num_splits]
-        kl = part * (np.log(part) - np.log(np.mean(part, axis=0, keepdims=True)))
-        kl = np.mean(np.sum(kl, axis=1))
-        scores.append(np.exp(kl))
-    return float(np.mean(scores)), float(np.std(scores))
-
-#----------------------------------------------------------------------------

metrics/kernel_inception_distance.py

@@ -1,46 +0,0 @@
-# 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.
-
-"""Kernel Inception Distance (KID) from the paper "Demystifying MMD
-GANs". Matches the original implementation by Binkowski et al. at
-https://github.com/mbinkowski/MMD-GAN/blob/master/gan/compute_scores.py"""
-
-import numpy as np
-from . import metric_utils
-
-#----------------------------------------------------------------------------
-
-def compute_kid(opts, max_real, num_gen, num_subsets, max_subset_size):
-    # Direct TorchScript translation of http://download.tensorflow.org/models/image/imagenet/inception-2015-12-05.tgz
-    detector_url = 'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metrics/inception-2015-12-05.pt'
-    detector_kwargs = dict(return_features=True) # Return raw features before the softmax layer.
-
-    real_features = metric_utils.compute_feature_stats_for_dataset(
-        opts=opts, detector_url=detector_url, detector_kwargs=detector_kwargs,
-        rel_lo=0, rel_hi=0, capture_all=True, max_items=max_real).get_all()
-
-    gen_features = metric_utils.compute_feature_stats_for_generator(
-        opts=opts, detector_url=detector_url, detector_kwargs=detector_kwargs,
-        rel_lo=0, rel_hi=1, capture_all=True, max_items=num_gen).get_all()
-
-    if opts.rank != 0:
-        return float('nan')
-
-    n = real_features.shape[1]
-    m = min(min(real_features.shape[0], gen_features.shape[0]), max_subset_size)
-    t = 0
-    for _subset_idx in range(num_subsets):
-        x = gen_features[np.random.choice(gen_features.shape[0], m, replace=False)]
-        y = real_features[np.random.choice(real_features.shape[0], m, replace=False)]
-        a = (x @ x.T / n + 1) ** 3 + (y @ y.T / n + 1) ** 3
-        b = (x @ y.T / n + 1) ** 3
-        t += (a.sum() - np.diag(a).sum()) / (m - 1) - b.sum() * 2 / m
-    kid = t / num_subsets / m
-    return float(kid)
-
-#----------------------------------------------------------------------------

metrics/metric_main.py

@@ -1,152 +0,0 @@
-# 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.
-
-import os
-import time
-import json
-import torch
-import dnnlib
-
-from . import metric_utils
-from . import frechet_inception_distance
-from . import kernel_inception_distance
-from . import precision_recall
-from . import perceptual_path_length
-from . import inception_score
-
-#----------------------------------------------------------------------------
-
-_metric_dict = dict() # name => fn
-
-def register_metric(fn):
-    assert callable(fn)
-    _metric_dict[fn.__name__] = fn
-    return fn
-
-def is_valid_metric(metric):
-    return metric in _metric_dict
-
-def list_valid_metrics():
-    return list(_metric_dict.keys())
-
-#----------------------------------------------------------------------------
-
-def calc_metric(metric, **kwargs): # See metric_utils.MetricOptions for the full list of arguments.
-    assert is_valid_metric(metric)
-    opts = metric_utils.MetricOptions(**kwargs)
-
-    # Calculate.
-    start_time = time.time()
-    results = _metric_dict[metric](opts)
-    total_time = time.time() - start_time
-
-    # Broadcast results.
-    for key, value in list(results.items()):
-        if opts.num_gpus > 1:
-            value = torch.as_tensor(value, dtype=torch.float64, device=opts.device)
-            torch.distributed.broadcast(tensor=value, src=0)
-            value = float(value.cpu())
-        results[key] = value
-
-    # Decorate with metadata.
-    return dnnlib.EasyDict(
-        results         = dnnlib.EasyDict(results),
-        metric          = metric,
-        total_time      = total_time,
-        total_time_str  = dnnlib.util.format_time(total_time),
-        num_gpus        = opts.num_gpus,
-    )
-
-#----------------------------------------------------------------------------
-
-def report_metric(result_dict, run_dir=None, snapshot_pkl=None):
-    metric = result_dict['metric']
-    assert is_valid_metric(metric)
-    if run_dir is not None and snapshot_pkl is not None:
-        snapshot_pkl = os.path.relpath(snapshot_pkl, run_dir)
-
-    jsonl_line = json.dumps(dict(result_dict, snapshot_pkl=snapshot_pkl, timestamp=time.time()))
-    print(jsonl_line)
-    if run_dir is not None and os.path.isdir(run_dir):
-        with open(os.path.join(run_dir, f'metric-{metric}.jsonl'), 'at') as f:
-            f.write(jsonl_line + '\n')
-
-#----------------------------------------------------------------------------
-# Primary metrics.
-
-@register_metric
-def fid50k_full(opts):
-    opts.dataset_kwargs.update(max_size=None, xflip=False)
-    fid = frechet_inception_distance.compute_fid(opts, max_real=None, num_gen=50000)
-    return dict(fid50k_full=fid)
-
-@register_metric
-def kid50k_full(opts):
-    opts.dataset_kwargs.update(max_size=None, xflip=False)
-    kid = kernel_inception_distance.compute_kid(opts, max_real=1000000, num_gen=50000, num_subsets=100, max_subset_size=1000)
-    return dict(kid50k_full=kid)
-
-@register_metric
-def pr50k3_full(opts):
-    opts.dataset_kwargs.update(max_size=None, xflip=False)
-    precision, recall = precision_recall.compute_pr(opts, max_real=200000, num_gen=50000, nhood_size=3, row_batch_size=10000, col_batch_size=10000)
-    return dict(pr50k3_full_precision=precision, pr50k3_full_recall=recall)
-
-@register_metric
-def ppl2_wend(opts):
-    ppl = perceptual_path_length.compute_ppl(opts, num_samples=50000, epsilon=1e-4, space='w', sampling='end', crop=False, batch_size=2)
-    return dict(ppl2_wend=ppl)
-
-@register_metric
-def is50k(opts):
-    opts.dataset_kwargs.update(max_size=None, xflip=False)
-    mean, std = inception_score.compute_is(opts, num_gen=50000, num_splits=10)
-    return dict(is50k_mean=mean, is50k_std=std)
-
-#----------------------------------------------------------------------------
-# Legacy metrics.
-
-@register_metric
-def fid50k(opts):
-    opts.dataset_kwargs.update(max_size=None)
-    fid = frechet_inception_distance.compute_fid(opts, max_real=50000, num_gen=50000)
-    return dict(fid50k=fid)
-
-@register_metric
-def kid50k(opts):
-    opts.dataset_kwargs.update(max_size=None)
-    kid = kernel_inception_distance.compute_kid(opts, max_real=50000, num_gen=50000, num_subsets=100, max_subset_size=1000)
-    return dict(kid50k=kid)
-
-@register_metric
-def pr50k3(opts):
-    opts.dataset_kwargs.update(max_size=None)
-    precision, recall = precision_recall.compute_pr(opts, max_real=50000, num_gen=50000, nhood_size=3, row_batch_size=10000, col_batch_size=10000)
-    return dict(pr50k3_precision=precision, pr50k3_recall=recall)
-
-@register_metric
-def ppl_zfull(opts):
-    ppl = perceptual_path_length.compute_ppl(opts, num_samples=50000, epsilon=1e-4, space='z', sampling='full', crop=True, batch_size=2)
-    return dict(ppl_zfull=ppl)
-
-@register_metric
-def ppl_wfull(opts):
-    ppl = perceptual_path_length.compute_ppl(opts, num_samples=50000, epsilon=1e-4, space='w', sampling='full', crop=True, batch_size=2)
-    return dict(ppl_wfull=ppl)
-
-@register_metric
-def ppl_zend(opts):
-    ppl = perceptual_path_length.compute_ppl(opts, num_samples=50000, epsilon=1e-4, space='z', sampling='end', crop=True, batch_size=2)
-    return dict(ppl_zend=ppl)
-
-@register_metric
-def ppl_wend(opts):
-    ppl = perceptual_path_length.compute_ppl(opts, num_samples=50000, epsilon=1e-4, space='w', sampling='end', crop=True, batch_size=2)
-    return dict(ppl_wend=ppl)
-
-#----------------------------------------------------------------------------

metrics/metric_utils.py

@@ -1,275 +0,0 @@
-# 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.
-
-import os
-import time
-import hashlib
-import pickle
-import copy
-import uuid
-import numpy as np
-import torch
-import dnnlib
-
-#----------------------------------------------------------------------------
-
-class MetricOptions:
-    def __init__(self, G=None, G_kwargs={}, dataset_kwargs={}, num_gpus=1, rank=0, device=None, progress=None, cache=True):
-        assert 0 <= rank < num_gpus
-        self.G              = G
-        self.G_kwargs       = dnnlib.EasyDict(G_kwargs)
-        self.dataset_kwargs = dnnlib.EasyDict(dataset_kwargs)
-        self.num_gpus       = num_gpus
-        self.rank           = rank
-        self.device         = device if device is not None else torch.device('cuda', rank)
-        self.progress       = progress.sub() if progress is not None and rank == 0 else ProgressMonitor()
-        self.cache          = cache
-
-#----------------------------------------------------------------------------
-
-_feature_detector_cache = dict()
-
-def get_feature_detector_name(url):
-    return os.path.splitext(url.split('/')[-1])[0]
-
-def get_feature_detector(url, device=torch.device('cpu'), num_gpus=1, rank=0, verbose=False):
-    assert 0 <= rank < num_gpus
-    key = (url, device)
-    if key not in _feature_detector_cache:
-        is_leader = (rank == 0)
-        if not is_leader and num_gpus > 1:
-            torch.distributed.barrier() # leader goes first
-        with dnnlib.util.open_url(url, verbose=(verbose and is_leader)) as f:
-            _feature_detector_cache[key] = torch.jit.load(f).eval().to(device)
-        if is_leader and num_gpus > 1:
-            torch.distributed.barrier() # others follow
-    return _feature_detector_cache[key]
-
-#----------------------------------------------------------------------------
-
-class FeatureStats:
-    def __init__(self, capture_all=False, capture_mean_cov=False, max_items=None):
-        self.capture_all = capture_all
-        self.capture_mean_cov = capture_mean_cov
-        self.max_items = max_items
-        self.num_items = 0
-        self.num_features = None
-        self.all_features = None
-        self.raw_mean = None
-        self.raw_cov = None
-
-    def set_num_features(self, num_features):
-        if self.num_features is not None:
-            assert num_features == self.num_features
-        else:
-            self.num_features = num_features
-            self.all_features = []
-            self.raw_mean = np.zeros([num_features], dtype=np.float64)
-            self.raw_cov = np.zeros([num_features, num_features], dtype=np.float64)
-
-    def is_full(self):
-        return (self.max_items is not None) and (self.num_items >= self.max_items)
-
-    def append(self, x):
-        x = np.asarray(x, dtype=np.float32)
-        assert x.ndim == 2
-        if (self.max_items is not None) and (self.num_items + x.shape[0] > self.max_items):
-            if self.num_items >= self.max_items:
-                return
-            x = x[:self.max_items - self.num_items]
-
-        self.set_num_features(x.shape[1])
-        self.num_items += x.shape[0]
-        if self.capture_all:
-            self.all_features.append(x)
-        if self.capture_mean_cov:
-            x64 = x.astype(np.float64)
-            self.raw_mean += x64.sum(axis=0)
-            self.raw_cov += x64.T @ x64
-
-    def append_torch(self, x, num_gpus=1, rank=0):
-        assert isinstance(x, torch.Tensor) and x.ndim == 2
-        assert 0 <= rank < num_gpus
-        if num_gpus > 1:
-            ys = []
-            for src in range(num_gpus):
-                y = x.clone()
-                torch.distributed.broadcast(y, src=src)
-                ys.append(y)
-            x = torch.stack(ys, dim=1).flatten(0, 1) # interleave samples
-        self.append(x.cpu().numpy())
-
-    def get_all(self):
-        assert self.capture_all
-        return np.concatenate(self.all_features, axis=0)
-
-    def get_all_torch(self):
-        return torch.from_numpy(self.get_all())
-
-    def get_mean_cov(self):
-        assert self.capture_mean_cov
-        mean = self.raw_mean / self.num_items
-        cov = self.raw_cov / self.num_items
-        cov = cov - np.outer(mean, mean)
-        return mean, cov
-
-    def save(self, pkl_file):
-        with open(pkl_file, 'wb') as f:
-            pickle.dump(self.__dict__, f)
-
-    @staticmethod
-    def load(pkl_file):
-        with open(pkl_file, 'rb') as f:
-            s = dnnlib.EasyDict(pickle.load(f))
-        obj = FeatureStats(capture_all=s.capture_all, max_items=s.max_items)
-        obj.__dict__.update(s)
-        return obj
-
-#----------------------------------------------------------------------------
-
-class ProgressMonitor:
-    def __init__(self, tag=None, num_items=None, flush_interval=1000, verbose=False, progress_fn=None, pfn_lo=0, pfn_hi=1000, pfn_total=1000):
-        self.tag = tag
-        self.num_items = num_items
-        self.verbose = verbose
-        self.flush_interval = flush_interval
-        self.progress_fn = progress_fn
-        self.pfn_lo = pfn_lo
-        self.pfn_hi = pfn_hi
-        self.pfn_total = pfn_total
-        self.start_time = time.time()
-        self.batch_time = self.start_time
-        self.batch_items = 0
-        if self.progress_fn is not None:
-            self.progress_fn(self.pfn_lo, self.pfn_total)
-
-    def update(self, cur_items):
-        assert (self.num_items is None) or (cur_items <= self.num_items)
-        if (cur_items < self.batch_items + self.flush_interval) and (self.num_items is None or cur_items < self.num_items):
-            return
-        cur_time = time.time()
-        total_time = cur_time - self.start_time
-        time_per_item = (cur_time - self.batch_time) / max(cur_items - self.batch_items, 1)
-        if (self.verbose) and (self.tag is not None):
-            print(f'{self.tag:<19s} items {cur_items:<7d} time {dnnlib.util.format_time(total_time):<12s} ms/item {time_per_item*1e3:.2f}')
-        self.batch_time = cur_time
-        self.batch_items = cur_items
-
-        if (self.progress_fn is not None) and (self.num_items is not None):
-            self.progress_fn(self.pfn_lo + (self.pfn_hi - self.pfn_lo) * (cur_items / self.num_items), self.pfn_total)
-
-    def sub(self, tag=None, num_items=None, flush_interval=1000, rel_lo=0, rel_hi=1):
-        return ProgressMonitor(
-            tag             = tag,
-            num_items       = num_items,
-            flush_interval  = flush_interval,
-            verbose         = self.verbose,
-            progress_fn     = self.progress_fn,
-            pfn_lo          = self.pfn_lo + (self.pfn_hi - self.pfn_lo) * rel_lo,
-            pfn_hi          = self.pfn_lo + (self.pfn_hi - self.pfn_lo) * rel_hi,
-            pfn_total       = self.pfn_total,
-        )
-
-#----------------------------------------------------------------------------
-
-def compute_feature_stats_for_dataset(opts, detector_url, detector_kwargs, rel_lo=0, rel_hi=1, batch_size=64, data_loader_kwargs=None, max_items=None, **stats_kwargs):
-    dataset = dnnlib.util.construct_class_by_name(**opts.dataset_kwargs)
-    if data_loader_kwargs is None:
-        data_loader_kwargs = dict(pin_memory=True, num_workers=3, prefetch_factor=2)
-
-    # Try to lookup from cache.
-    cache_file = None
-    if opts.cache:
-        # Choose cache file name.
-        args = dict(dataset_kwargs=opts.dataset_kwargs, detector_url=detector_url, detector_kwargs=detector_kwargs, stats_kwargs=stats_kwargs)
-        md5 = hashlib.md5(repr(sorted(args.items())).encode('utf-8'))
-        cache_tag = f'{dataset.name}-{get_feature_detector_name(detector_url)}-{md5.hexdigest()}'
-        cache_file = dnnlib.make_cache_dir_path('gan-metrics', cache_tag + '.pkl')
-
-        # Check if the file exists (all processes must agree).
-        flag = os.path.isfile(cache_file) if opts.rank == 0 else False
-        if opts.num_gpus > 1:
-            flag = torch.as_tensor(flag, dtype=torch.float32, device=opts.device)
-            torch.distributed.broadcast(tensor=flag, src=0)
-            flag = (float(flag.cpu()) != 0)
-
-        # Load.
-        if flag:
-            return FeatureStats.load(cache_file)
-
-    # Initialize.
-    num_items = len(dataset)
-    if max_items is not None:
-        num_items = min(num_items, max_items)
-    stats = FeatureStats(max_items=num_items, **stats_kwargs)
-    progress = opts.progress.sub(tag='dataset features', num_items=num_items, rel_lo=rel_lo, rel_hi=rel_hi)
-    detector = get_feature_detector(url=detector_url, device=opts.device, num_gpus=opts.num_gpus, rank=opts.rank, verbose=progress.verbose)
-
-    # Main loop.
-    item_subset = [(i * opts.num_gpus + opts.rank) % num_items for i in range((num_items - 1) // opts.num_gpus + 1)]
-    for images, _labels in torch.utils.data.DataLoader(dataset=dataset, sampler=item_subset, batch_size=batch_size, **data_loader_kwargs):
-        if images.shape[1] == 1:
-            images = images.repeat([1, 3, 1, 1])
-        features = detector(images.to(opts.device), **detector_kwargs)
-        stats.append_torch(features, num_gpus=opts.num_gpus, rank=opts.rank)
-        progress.update(stats.num_items)
-
-    # Save to cache.
-    if cache_file is not None and opts.rank == 0:
-        os.makedirs(os.path.dirname(cache_file), exist_ok=True)
-        temp_file = cache_file + '.' + uuid.uuid4().hex
-        stats.save(temp_file)
-        os.replace(temp_file, cache_file) # atomic
-    return stats
-
-#----------------------------------------------------------------------------
-
-def compute_feature_stats_for_generator(opts, detector_url, detector_kwargs, rel_lo=0, rel_hi=1, batch_size=64, batch_gen=None, jit=False, **stats_kwargs):
-    if batch_gen is None:
-        batch_gen = min(batch_size, 4)
-    assert batch_size % batch_gen == 0
-
-    # Setup generator and load labels.
-    G = copy.deepcopy(opts.G).eval().requires_grad_(False).to(opts.device)
-    dataset = dnnlib.util.construct_class_by_name(**opts.dataset_kwargs)
-
-    # Image generation func.
-    def run_generator(z, c):
-        img = G(z=z, c=c, **opts.G_kwargs)
-        img = (img * 127.5 + 128).clamp(0, 255).to(torch.uint8)
-        return img
-
-    # JIT.
-    if jit:
-        z = torch.zeros([batch_gen, G.z_dim], device=opts.device)
-        c = torch.zeros([batch_gen, G.c_dim], device=opts.device)
-        run_generator = torch.jit.trace(run_generator, [z, c], check_trace=False)
-
-    # Initialize.
-    stats = FeatureStats(**stats_kwargs)
-    assert stats.max_items is not None
-    progress = opts.progress.sub(tag='generator features', num_items=stats.max_items, rel_lo=rel_lo, rel_hi=rel_hi)
-    detector = get_feature_detector(url=detector_url, device=opts.device, num_gpus=opts.num_gpus, rank=opts.rank, verbose=progress.verbose)
-
-    # Main loop.
-    while not stats.is_full():
-        images = []
-        for _i in range(batch_size // batch_gen):
-            z = torch.randn([batch_gen, G.z_dim], device=opts.device)
-            c = [dataset.get_label(np.random.randint(len(dataset))) for _i in range(batch_gen)]
-            c = torch.from_numpy(np.stack(c)).pin_memory().to(opts.device)
-            images.append(run_generator(z, c))
-        images = torch.cat(images)
-        if images.shape[1] == 1:
-            images = images.repeat([1, 3, 1, 1])
-        features = detector(images, **detector_kwargs)
-        stats.append_torch(features, num_gpus=opts.num_gpus, rank=opts.rank)
-        progress.update(stats.num_items)
-    return stats
-
-#----------------------------------------------------------------------------

metrics/perceptual_path_length.py

@@ -1,131 +0,0 @@
-# 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.
-
-"""Perceptual Path Length (PPL) from the paper "A Style-Based Generator
-Architecture for Generative Adversarial Networks". Matches the original
-implementation by Karras et al. at
-https://github.com/NVlabs/stylegan/blob/master/metrics/perceptual_path_length.py"""
-
-import copy
-import numpy as np
-import torch
-import dnnlib
-from . import metric_utils
-
-#----------------------------------------------------------------------------
-
-# Spherical interpolation of a batch of vectors.
-def slerp(a, b, t):
-    a = a / a.norm(dim=-1, keepdim=True)
-    b = b / b.norm(dim=-1, keepdim=True)
-    d = (a * b).sum(dim=-1, keepdim=True)
-    p = t * torch.acos(d)
-    c = b - d * a
-    c = c / c.norm(dim=-1, keepdim=True)
-    d = a * torch.cos(p) + c * torch.sin(p)
-    d = d / d.norm(dim=-1, keepdim=True)
-    return d
-
-#----------------------------------------------------------------------------
-
-class PPLSampler(torch.nn.Module):
-    def __init__(self, G, G_kwargs, epsilon, space, sampling, crop, vgg16):
-        assert space in ['z', 'w']
-        assert sampling in ['full', 'end']
-        super().__init__()
-        self.G = copy.deepcopy(G)
-        self.G_kwargs = G_kwargs
-        self.epsilon = epsilon
-        self.space = space
-        self.sampling = sampling
-        self.crop = crop
-        self.vgg16 = copy.deepcopy(vgg16)
-
-    def forward(self, c):
-        # Generate random latents and interpolation t-values.
-        t = torch.rand([c.shape[0]], device=c.device) * (1 if self.sampling == 'full' else 0)
-        z0, z1 = torch.randn([c.shape[0] * 2, self.G.z_dim], device=c.device).chunk(2)
-
-        # Interpolate in W or Z.
-        if self.space == 'w':
-            w0, w1 = self.G.mapping(z=torch.cat([z0,z1]), c=torch.cat([c,c])).chunk(2)
-            wt0 = w0.lerp(w1, t.unsqueeze(1).unsqueeze(2))
-            wt1 = w0.lerp(w1, t.unsqueeze(1).unsqueeze(2) + self.epsilon)
-        else: # space == 'z'
-            zt0 = slerp(z0, z1, t.unsqueeze(1))
-            zt1 = slerp(z0, z1, t.unsqueeze(1) + self.epsilon)
-            wt0, wt1 = self.G.mapping(z=torch.cat([zt0,zt1]), c=torch.cat([c,c])).chunk(2)
-
-        # Randomize noise buffers.
-        for name, buf in self.G.named_buffers():
-            if name.endswith('.noise_const'):
-                buf.copy_(torch.randn_like(buf))
-
-        # Generate images.
-        img = self.G.synthesis(ws=torch.cat([wt0,wt1]), noise_mode='const', force_fp32=True, **self.G_kwargs)
-
-        # Center crop.
-        if self.crop:
-            assert img.shape[2] == img.shape[3]
-            c = img.shape[2] // 8
-            img = img[:, :, c*3 : c*7, c*2 : c*6]
-
-        # Downsample to 256x256.
-        factor = self.G.img_resolution // 256
-        if factor > 1:
-            img = img.reshape([-1, img.shape[1], img.shape[2] // factor, factor, img.shape[3] // factor, factor]).mean([3, 5])
-
-        # Scale dynamic range from [-1,1] to [0,255].
-        img = (img + 1) * (255 / 2)
-        if self.G.img_channels == 1:
-            img = img.repeat([1, 3, 1, 1])
-
-        # Evaluate differential LPIPS.
-        lpips_t0, lpips_t1 = self.vgg16(img, resize_images=False, return_lpips=True).chunk(2)
-        dist = (lpips_t0 - lpips_t1).square().sum(1) / self.epsilon ** 2
-        return dist
-
-#----------------------------------------------------------------------------
-
-def compute_ppl(opts, num_samples, epsilon, space, sampling, crop, batch_size, jit=False):
-    dataset = dnnlib.util.construct_class_by_name(**opts.dataset_kwargs)
-    vgg16_url = 'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metrics/vgg16.pt'
-    vgg16 = metric_utils.get_feature_detector(vgg16_url, num_gpus=opts.num_gpus, rank=opts.rank, verbose=opts.progress.verbose)
-
-    # Setup sampler.
-    sampler = PPLSampler(G=opts.G, G_kwargs=opts.G_kwargs, epsilon=epsilon, space=space, sampling=sampling, crop=crop, vgg16=vgg16)
-    sampler.eval().requires_grad_(False).to(opts.device)
-    if jit:
-        c = torch.zeros([batch_size, opts.G.c_dim], device=opts.device)
-        sampler = torch.jit.trace(sampler, [c], check_trace=False)
-
-    # Sampling loop.
-    dist = []
-    progress = opts.progress.sub(tag='ppl sampling', num_items=num_samples)
-    for batch_start in range(0, num_samples, batch_size * opts.num_gpus):
-        progress.update(batch_start)
-        c = [dataset.get_label(np.random.randint(len(dataset))) for _i in range(batch_size)]
-        c = torch.from_numpy(np.stack(c)).pin_memory().to(opts.device)
-        x = sampler(c)
-        for src in range(opts.num_gpus):
-            y = x.clone()
-            if opts.num_gpus > 1:
-                torch.distributed.broadcast(y, src=src)
-            dist.append(y)
-    progress.update(num_samples)
-
-    # Compute PPL.
-    if opts.rank != 0:
-        return float('nan')
-    dist = torch.cat(dist)[:num_samples].cpu().numpy()
-    lo = np.percentile(dist, 1, interpolation='lower')
-    hi = np.percentile(dist, 99, interpolation='higher')
-    ppl = np.extract(np.logical_and(dist >= lo, dist <= hi), dist).mean()
-    return float(ppl)
-
-#----------------------------------------------------------------------------

metrics/precision_recall.py

@@ -1,62 +0,0 @@
-# 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.
-
-"""Precision/Recall (PR) from the paper "Improved Precision and Recall
-Metric for Assessing Generative Models". Matches the original implementation
-by Kynkaanniemi et al. at
-https://github.com/kynkaat/improved-precision-and-recall-metric/blob/master/precision_recall.py"""
-
-import torch
-from . import metric_utils
-
-#----------------------------------------------------------------------------
-
-def compute_distances(row_features, col_features, num_gpus, rank, col_batch_size):
-    assert 0 <= rank < num_gpus
-    num_cols = col_features.shape[0]
-    num_batches = ((num_cols - 1) // col_batch_size // num_gpus + 1) * num_gpus
-    col_batches = torch.nn.functional.pad(col_features, [0, 0, 0, -num_cols % num_batches]).chunk(num_batches)
-    dist_batches = []
-    for col_batch in col_batches[rank :: num_gpus]:
-        dist_batch = torch.cdist(row_features.unsqueeze(0), col_batch.unsqueeze(0))[0]
-        for src in range(num_gpus):
-            dist_broadcast = dist_batch.clone()
-            if num_gpus > 1:
-                torch.distributed.broadcast(dist_broadcast, src=src)
-            dist_batches.append(dist_broadcast.cpu() if rank == 0 else None)
-    return torch.cat(dist_batches, dim=1)[:, :num_cols] if rank == 0 else None
-
-#----------------------------------------------------------------------------
-
-def compute_pr(opts, max_real, num_gen, nhood_size, row_batch_size, col_batch_size):
-    detector_url = 'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metrics/vgg16.pt'
-    detector_kwargs = dict(return_features=True)
-
-    real_features = metric_utils.compute_feature_stats_for_dataset(
-        opts=opts, detector_url=detector_url, detector_kwargs=detector_kwargs,
-        rel_lo=0, rel_hi=0, capture_all=True, max_items=max_real).get_all_torch().to(torch.float16).to(opts.device)
-
-    gen_features = metric_utils.compute_feature_stats_for_generator(
-        opts=opts, detector_url=detector_url, detector_kwargs=detector_kwargs,
-        rel_lo=0, rel_hi=1, capture_all=True, max_items=num_gen).get_all_torch().to(torch.float16).to(opts.device)
-
-    results = dict()
-    for name, manifold, probes in [('precision', real_features, gen_features), ('recall', gen_features, real_features)]:
-        kth = []
-        for manifold_batch in manifold.split(row_batch_size):
-            dist = compute_distances(row_features=manifold_batch, col_features=manifold, num_gpus=opts.num_gpus, rank=opts.rank, col_batch_size=col_batch_size)
-            kth.append(dist.to(torch.float32).kthvalue(nhood_size + 1).values.to(torch.float16) if opts.rank == 0 else None)
-        kth = torch.cat(kth) if opts.rank == 0 else None
-        pred = []
-        for probes_batch in probes.split(row_batch_size):
-            dist = compute_distances(row_features=probes_batch, col_features=manifold, num_gpus=opts.num_gpus, rank=opts.rank, col_batch_size=col_batch_size)
-            pred.append((dist <= kth).any(dim=1) if opts.rank == 0 else None)
-        results[name] = float(torch.cat(pred).to(torch.float32).mean() if opts.rank == 0 else 'nan')
-    return results['precision'], results['recall']
-
-#----------------------------------------------------------------------------

pretrained/ffhq.pkl

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

pretrained/metfaces.pkl

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