update
This commit is contained in:
chenxuanhong
@@ -0,0 +1,267 @@
|
||||
# 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.
|
||||
|
||||
"""Equivariance metrics (EQ-T, EQ-T_frac, and EQ-R) from the paper
|
||||
"Alias-Free Generative Adversarial Networks"."""
|
||||
|
||||
import copy
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.fft
|
||||
from torch_utils.ops import upfirdn2d
|
||||
from . import metric_utils
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
# Utilities.
|
||||
|
||||
def sinc(x):
|
||||
y = (x * np.pi).abs()
|
||||
z = torch.sin(y) / y.clamp(1e-30, float('inf'))
|
||||
return torch.where(y < 1e-30, torch.ones_like(x), z)
|
||||
|
||||
def lanczos_window(x, a):
|
||||
x = x.abs() / a
|
||||
return torch.where(x < 1, sinc(x), torch.zeros_like(x))
|
||||
|
||||
def rotation_matrix(angle):
|
||||
angle = torch.as_tensor(angle).to(torch.float32)
|
||||
mat = torch.eye(3, device=angle.device)
|
||||
mat[0, 0] = angle.cos()
|
||||
mat[0, 1] = angle.sin()
|
||||
mat[1, 0] = -angle.sin()
|
||||
mat[1, 1] = angle.cos()
|
||||
return mat
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
# Apply integer translation to a batch of 2D images. Corresponds to the
|
||||
# operator T_x in Appendix E.1.
|
||||
|
||||
def apply_integer_translation(x, tx, ty):
|
||||
_N, _C, H, W = x.shape
|
||||
tx = torch.as_tensor(tx * W).to(dtype=torch.float32, device=x.device)
|
||||
ty = torch.as_tensor(ty * H).to(dtype=torch.float32, device=x.device)
|
||||
ix = tx.round().to(torch.int64)
|
||||
iy = ty.round().to(torch.int64)
|
||||
|
||||
z = torch.zeros_like(x)
|
||||
m = torch.zeros_like(x)
|
||||
if abs(ix) < W and abs(iy) < H:
|
||||
y = x[:, :, max(-iy,0) : H+min(-iy,0), max(-ix,0) : W+min(-ix,0)]
|
||||
z[:, :, max(iy,0) : H+min(iy,0), max(ix,0) : W+min(ix,0)] = y
|
||||
m[:, :, max(iy,0) : H+min(iy,0), max(ix,0) : W+min(ix,0)] = 1
|
||||
return z, m
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
# Apply integer translation to a batch of 2D images. Corresponds to the
|
||||
# operator T_x in Appendix E.2.
|
||||
|
||||
def apply_fractional_translation(x, tx, ty, a=3):
|
||||
_N, _C, H, W = x.shape
|
||||
tx = torch.as_tensor(tx * W).to(dtype=torch.float32, device=x.device)
|
||||
ty = torch.as_tensor(ty * H).to(dtype=torch.float32, device=x.device)
|
||||
ix = tx.floor().to(torch.int64)
|
||||
iy = ty.floor().to(torch.int64)
|
||||
fx = tx - ix
|
||||
fy = ty - iy
|
||||
b = a - 1
|
||||
|
||||
z = torch.zeros_like(x)
|
||||
zx0 = max(ix - b, 0)
|
||||
zy0 = max(iy - b, 0)
|
||||
zx1 = min(ix + a, 0) + W
|
||||
zy1 = min(iy + a, 0) + H
|
||||
if zx0 < zx1 and zy0 < zy1:
|
||||
taps = torch.arange(a * 2, device=x.device) - b
|
||||
filter_x = (sinc(taps - fx) * sinc((taps - fx) / a)).unsqueeze(0)
|
||||
filter_y = (sinc(taps - fy) * sinc((taps - fy) / a)).unsqueeze(1)
|
||||
y = x
|
||||
y = upfirdn2d.filter2d(y, filter_x / filter_x.sum(), padding=[b,a,0,0])
|
||||
y = upfirdn2d.filter2d(y, filter_y / filter_y.sum(), padding=[0,0,b,a])
|
||||
y = y[:, :, max(b-iy,0) : H+b+a+min(-iy-a,0), max(b-ix,0) : W+b+a+min(-ix-a,0)]
|
||||
z[:, :, zy0:zy1, zx0:zx1] = y
|
||||
|
||||
m = torch.zeros_like(x)
|
||||
mx0 = max(ix + a, 0)
|
||||
my0 = max(iy + a, 0)
|
||||
mx1 = min(ix - b, 0) + W
|
||||
my1 = min(iy - b, 0) + H
|
||||
if mx0 < mx1 and my0 < my1:
|
||||
m[:, :, my0:my1, mx0:mx1] = 1
|
||||
return z, m
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
# Construct an oriented low-pass filter that applies the appropriate
|
||||
# bandlimit with respect to the input and output of the given affine 2D
|
||||
# image transformation.
|
||||
|
||||
def construct_affine_bandlimit_filter(mat, a=3, amax=16, aflt=64, up=4, cutoff_in=1, cutoff_out=1):
|
||||
assert a <= amax < aflt
|
||||
mat = torch.as_tensor(mat).to(torch.float32)
|
||||
|
||||
# Construct 2D filter taps in input & output coordinate spaces.
|
||||
taps = ((torch.arange(aflt * up * 2 - 1, device=mat.device) + 1) / up - aflt).roll(1 - aflt * up)
|
||||
yi, xi = torch.meshgrid(taps, taps)
|
||||
xo, yo = (torch.stack([xi, yi], dim=2) @ mat[:2, :2].t()).unbind(2)
|
||||
|
||||
# Convolution of two oriented 2D sinc filters.
|
||||
fi = sinc(xi * cutoff_in) * sinc(yi * cutoff_in)
|
||||
fo = sinc(xo * cutoff_out) * sinc(yo * cutoff_out)
|
||||
f = torch.fft.ifftn(torch.fft.fftn(fi) * torch.fft.fftn(fo)).real
|
||||
|
||||
# Convolution of two oriented 2D Lanczos windows.
|
||||
wi = lanczos_window(xi, a) * lanczos_window(yi, a)
|
||||
wo = lanczos_window(xo, a) * lanczos_window(yo, a)
|
||||
w = torch.fft.ifftn(torch.fft.fftn(wi) * torch.fft.fftn(wo)).real
|
||||
|
||||
# Construct windowed FIR filter.
|
||||
f = f * w
|
||||
|
||||
# Finalize.
|
||||
c = (aflt - amax) * up
|
||||
f = f.roll([aflt * up - 1] * 2, dims=[0,1])[c:-c, c:-c]
|
||||
f = torch.nn.functional.pad(f, [0, 1, 0, 1]).reshape(amax * 2, up, amax * 2, up)
|
||||
f = f / f.sum([0,2], keepdim=True) / (up ** 2)
|
||||
f = f.reshape(amax * 2 * up, amax * 2 * up)[:-1, :-1]
|
||||
return f
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
# Apply the given affine transformation to a batch of 2D images.
|
||||
|
||||
def apply_affine_transformation(x, mat, up=4, **filter_kwargs):
|
||||
_N, _C, H, W = x.shape
|
||||
mat = torch.as_tensor(mat).to(dtype=torch.float32, device=x.device)
|
||||
|
||||
# Construct filter.
|
||||
f = construct_affine_bandlimit_filter(mat, up=up, **filter_kwargs)
|
||||
assert f.ndim == 2 and f.shape[0] == f.shape[1] and f.shape[0] % 2 == 1
|
||||
p = f.shape[0] // 2
|
||||
|
||||
# Construct sampling grid.
|
||||
theta = mat.inverse()
|
||||
theta[:2, 2] *= 2
|
||||
theta[0, 2] += 1 / up / W
|
||||
theta[1, 2] += 1 / up / H
|
||||
theta[0, :] *= W / (W + p / up * 2)
|
||||
theta[1, :] *= H / (H + p / up * 2)
|
||||
theta = theta[:2, :3].unsqueeze(0).repeat([x.shape[0], 1, 1])
|
||||
g = torch.nn.functional.affine_grid(theta, x.shape, align_corners=False)
|
||||
|
||||
# Resample image.
|
||||
y = upfirdn2d.upsample2d(x=x, f=f, up=up, padding=p)
|
||||
z = torch.nn.functional.grid_sample(y, g, mode='bilinear', padding_mode='zeros', align_corners=False)
|
||||
|
||||
# Form mask.
|
||||
m = torch.zeros_like(y)
|
||||
c = p * 2 + 1
|
||||
m[:, :, c:-c, c:-c] = 1
|
||||
m = torch.nn.functional.grid_sample(m, g, mode='nearest', padding_mode='zeros', align_corners=False)
|
||||
return z, m
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
# Apply fractional rotation to a batch of 2D images. Corresponds to the
|
||||
# operator R_\alpha in Appendix E.3.
|
||||
|
||||
def apply_fractional_rotation(x, angle, a=3, **filter_kwargs):
|
||||
angle = torch.as_tensor(angle).to(dtype=torch.float32, device=x.device)
|
||||
mat = rotation_matrix(angle)
|
||||
return apply_affine_transformation(x, mat, a=a, amax=a*2, **filter_kwargs)
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
# Modify the frequency content of a batch of 2D images as if they had undergo
|
||||
# fractional rotation -- but without actually rotating them. Corresponds to
|
||||
# the operator R^*_\alpha in Appendix E.3.
|
||||
|
||||
def apply_fractional_pseudo_rotation(x, angle, a=3, **filter_kwargs):
|
||||
angle = torch.as_tensor(angle).to(dtype=torch.float32, device=x.device)
|
||||
mat = rotation_matrix(-angle)
|
||||
f = construct_affine_bandlimit_filter(mat, a=a, amax=a*2, up=1, **filter_kwargs)
|
||||
y = upfirdn2d.filter2d(x=x, f=f)
|
||||
m = torch.zeros_like(y)
|
||||
c = f.shape[0] // 2
|
||||
m[:, :, c:-c, c:-c] = 1
|
||||
return y, m
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
# Compute the selected equivariance metrics for the given generator.
|
||||
|
||||
def compute_equivariance_metrics(opts, num_samples, batch_size, translate_max=0.125, rotate_max=1, compute_eqt_int=False, compute_eqt_frac=False, compute_eqr=False):
|
||||
assert compute_eqt_int or compute_eqt_frac or compute_eqr
|
||||
|
||||
# Setup generator and labels.
|
||||
G = copy.deepcopy(opts.G).eval().requires_grad_(False).to(opts.device)
|
||||
I = torch.eye(3, device=opts.device)
|
||||
M = getattr(getattr(getattr(G, 'synthesis', None), 'input', None), 'transform', None)
|
||||
if M is None:
|
||||
raise ValueError('Cannot compute equivariance metrics; the given generator does not support user-specified image transformations')
|
||||
c_iter = metric_utils.iterate_random_labels(opts=opts, batch_size=batch_size)
|
||||
|
||||
# Sampling loop.
|
||||
sums = None
|
||||
progress = opts.progress.sub(tag='eq sampling', num_items=num_samples)
|
||||
for batch_start in range(0, num_samples, batch_size * opts.num_gpus):
|
||||
progress.update(batch_start)
|
||||
s = []
|
||||
|
||||
# Randomize noise buffers, if any.
|
||||
for name, buf in G.named_buffers():
|
||||
if name.endswith('.noise_const'):
|
||||
buf.copy_(torch.randn_like(buf))
|
||||
|
||||
# Run mapping network.
|
||||
z = torch.randn([batch_size, G.z_dim], device=opts.device)
|
||||
c = next(c_iter)
|
||||
ws = G.mapping(z=z, c=c)
|
||||
|
||||
# Generate reference image.
|
||||
M[:] = I
|
||||
orig = G.synthesis(ws=ws, noise_mode='const', **opts.G_kwargs)
|
||||
|
||||
# Integer translation (EQ-T).
|
||||
if compute_eqt_int:
|
||||
t = (torch.rand(2, device=opts.device) * 2 - 1) * translate_max
|
||||
t = (t * G.img_resolution).round() / G.img_resolution
|
||||
M[:] = I
|
||||
M[:2, 2] = -t
|
||||
img = G.synthesis(ws=ws, noise_mode='const', **opts.G_kwargs)
|
||||
ref, mask = apply_integer_translation(orig, t[0], t[1])
|
||||
s += [(ref - img).square() * mask, mask]
|
||||
|
||||
# Fractional translation (EQ-T_frac).
|
||||
if compute_eqt_frac:
|
||||
t = (torch.rand(2, device=opts.device) * 2 - 1) * translate_max
|
||||
M[:] = I
|
||||
M[:2, 2] = -t
|
||||
img = G.synthesis(ws=ws, noise_mode='const', **opts.G_kwargs)
|
||||
ref, mask = apply_fractional_translation(orig, t[0], t[1])
|
||||
s += [(ref - img).square() * mask, mask]
|
||||
|
||||
# Rotation (EQ-R).
|
||||
if compute_eqr:
|
||||
angle = (torch.rand([], device=opts.device) * 2 - 1) * (rotate_max * np.pi)
|
||||
M[:] = rotation_matrix(-angle)
|
||||
img = G.synthesis(ws=ws, noise_mode='const', **opts.G_kwargs)
|
||||
ref, ref_mask = apply_fractional_rotation(orig, angle)
|
||||
pseudo, pseudo_mask = apply_fractional_pseudo_rotation(img, angle)
|
||||
mask = ref_mask * pseudo_mask
|
||||
s += [(ref - pseudo).square() * mask, mask]
|
||||
|
||||
# Accumulate results.
|
||||
s = torch.stack([x.to(torch.float64).sum() for x in s])
|
||||
sums = sums + s if sums is not None else s
|
||||
progress.update(num_samples)
|
||||
|
||||
# Compute PSNRs.
|
||||
if opts.num_gpus > 1:
|
||||
torch.distributed.all_reduce(sums)
|
||||
sums = sums.cpu()
|
||||
mses = sums[0::2] / sums[1::2]
|
||||
psnrs = np.log10(2) * 20 - mses.log10() * 10
|
||||
psnrs = tuple(psnrs.numpy())
|
||||
return psnrs[0] if len(psnrs) == 1 else psnrs
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
@@ -0,0 +1,41 @@
|
||||
# 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.
|
||||
|
||||
"""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, swav=False, sfid=False):
|
||||
# Direct TorchScript translation of http://download.tensorflow.org/models/image/imagenet/inception-2015-12-05.tgz
|
||||
detector_url = 'https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/metrics/inception-2015-12-05.pkl'
|
||||
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, swav=swav, sfid=sfid).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, swav=swav, sfid=sfid).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)
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
@@ -0,0 +1,38 @@
|
||||
# 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.
|
||||
|
||||
"""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://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/metrics/inception-2015-12-05.pkl'
|
||||
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))
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
@@ -0,0 +1,46 @@
|
||||
# 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.
|
||||
|
||||
"""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://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/metrics/inception-2015-12-05.pkl'
|
||||
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)
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
@@ -0,0 +1,151 @@
|
||||
# distribution of this software and related documentation without an express
|
||||
# license agreement from NVIDIA CORPORATION is strictly prohibited.
|
||||
|
||||
"""Main API for computing and reporting quality metrics."""
|
||||
|
||||
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
|
||||
from . import equivariance
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
_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')
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
# Recommended 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 fid10k_full(opts):
|
||||
opts.dataset_kwargs.update(max_size=None, xflip=False)
|
||||
fid = frechet_inception_distance.compute_fid(opts, max_real=None, num_gen=10000)
|
||||
return dict(fid10k_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 eqt50k_int(opts):
|
||||
opts.G_kwargs.update(force_fp32=True)
|
||||
psnr = equivariance.compute_equivariance_metrics(opts, num_samples=50000, batch_size=4, compute_eqt_int=True)
|
||||
return dict(eqt50k_int=psnr)
|
||||
|
||||
@register_metric
|
||||
def eqt50k_frac(opts):
|
||||
opts.G_kwargs.update(force_fp32=True)
|
||||
psnr = equivariance.compute_equivariance_metrics(opts, num_samples=50000, batch_size=4, compute_eqt_frac=True)
|
||||
return dict(eqt50k_frac=psnr)
|
||||
|
||||
@register_metric
|
||||
def eqr50k(opts):
|
||||
opts.G_kwargs.update(force_fp32=True)
|
||||
psnr = equivariance.compute_equivariance_metrics(opts, num_samples=50000, batch_size=4, compute_eqr=True)
|
||||
return dict(eqr50k=psnr)
|
||||
|
||||
# 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 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)
|
||||
@@ -0,0 +1,298 @@
|
||||
# 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.
|
||||
|
||||
"""Miscellaneous utilities used internally by the quality metrics."""
|
||||
|
||||
import os
|
||||
import time
|
||||
import hashlib
|
||||
import pickle
|
||||
import copy
|
||||
import uuid
|
||||
import numpy as np
|
||||
import torch
|
||||
import dnnlib
|
||||
from tqdm import tqdm
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
class MetricOptions:
|
||||
def __init__(self, G=None, G_kwargs={}, dataset_kwargs={}, num_gpus=1, rank=0, device=None, progress=None, cache=True, run_dir=None, cur_nimg=None, snapshot_pkl=None):
|
||||
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
|
||||
self.run_dir = run_dir
|
||||
self.cur_nimg = cur_nimg
|
||||
self.snapshot_pkl = snapshot_pkl
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
_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] = pickle.load(f).to(device)
|
||||
if is_leader and num_gpus > 1:
|
||||
torch.distributed.barrier() # others follow
|
||||
return _feature_detector_cache[key]
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
def iterate_random_labels(opts, batch_size):
|
||||
if opts.G.c_dim == 0:
|
||||
c = torch.zeros([batch_size, opts.G.c_dim], device=opts.device)
|
||||
while True:
|
||||
yield c
|
||||
else:
|
||||
dataset = dnnlib.util.construct_class_by_name(**opts.dataset_kwargs)
|
||||
while True:
|
||||
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)
|
||||
yield c
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
|
||||
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, swav=False, sfid=False, **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:
|
||||
det_name = get_feature_detector_name(detector_url)
|
||||
|
||||
# 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}-{det_name}-{md5.hexdigest()}'
|
||||
cache_file = os.path.join('.', 'dnnlib', 'gan-metrics', cache_tag + '.pkl')
|
||||
# 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)
|
||||
|
||||
print('Calculating the stats for this dataset the first time\n')
|
||||
print(f'Saving them to {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)
|
||||
|
||||
# get detector
|
||||
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 tqdm(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])
|
||||
|
||||
with torch.no_grad():
|
||||
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, swav=False, sfid=False, **stats_kwargs):
|
||||
if batch_gen is None:
|
||||
batch_gen = min(batch_size, 4)
|
||||
assert batch_size % batch_gen == 0
|
||||
|
||||
# Setup generator and labels.
|
||||
G = copy.deepcopy(opts.G).eval().requires_grad_(False).to(opts.device)
|
||||
c_iter = iterate_random_labels(opts=opts, batch_size=batch_gen)
|
||||
|
||||
# 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)
|
||||
|
||||
# get detector
|
||||
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)
|
||||
# img = G(z=z, c=next(c_iter), truncation_psi=0.1, **opts.G_kwargs)
|
||||
img = G(z=z, c=next(c_iter), **opts.G_kwargs)
|
||||
img = (img * 127.5 + 128).clamp(0, 255).to(torch.uint8)
|
||||
images.append(img)
|
||||
images = torch.cat(images)
|
||||
if images.shape[1] == 1:
|
||||
images = images.repeat([1, 3, 1, 1])
|
||||
|
||||
with torch.no_grad():
|
||||
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)
|
||||
return stats
|
||||
@@ -0,0 +1,125 @@
|
||||
# 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.
|
||||
|
||||
"""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
|
||||
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):
|
||||
vgg16_url = 'https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/metrics/vgg16.pkl'
|
||||
vgg16 = metric_utils.get_feature_detector(vgg16_url, num_gpus=opts.num_gpus, rank=opts.rank, verbose=opts.progress.verbose)
|
||||
|
||||
# Setup sampler and labels.
|
||||
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)
|
||||
c_iter = metric_utils.iterate_random_labels(opts=opts, batch_size=batch_size)
|
||||
|
||||
# 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)
|
||||
x = sampler(next(c_iter))
|
||||
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)
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
@@ -0,0 +1,62 @@
|
||||
# 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.
|
||||
|
||||
"""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://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/metrics/vgg16.pkl'
|
||||
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']
|
||||
|
||||
#----------------------------------------------------------------------------
|
||||
Reference in New Issue
Block a user