def glob_samples_paths(path, samples_find_deep, samples_find_ext, samples_ext_lossy=None, verbose=True):
vassert(type(samples_find_ext) is str and samples_find_ext != '', 'Sample extensions not specified')
vassert(
samples_ext_lossy is None or type(samples_ext_lossy) is str, 'Lossy sample extensions can be None or string'
)
vprint(verbose, f'Looking for samples {"recursively" if samples_find_deep else "non-recursivelty"} in "{path}" '
f'with extensions {samples_find_ext}')
samples_find_ext = [a.strip() for a in samples_find_ext.split(',') if a.strip() != '']
if samples_ext_lossy is not None:
samples_ext_lossy = [a.strip() for a in samples_ext_lossy.split(',') if a.strip() != '']
have_lossy = False
files = []
for r, d, ff in os.walk(path):
if not samples_find_deep and os.path.realpath(r) != os.path.realpath(path):
continue
for f in ff:
ext = os.path.splitext(f)[1].lower()
if len(ext) > 0 and ext[0] == '.':
ext = ext[1:]
if ext not in samples_find_ext:
continue
if samples_ext_lossy is not None and ext in samples_ext_lossy:
have_lossy = True
files.append(os.path.realpath(os.path.join(r, f)))
files = sorted(files)
vprint(verbose, f'Found {len(files)} samples'
f'{", some are lossy-compressed - this may affect metrics" if have_lossy else ""}')
return files
def create_feature_extractor(name, list_features, cuda=True, **kwargs):
vassert(name in FEATURE_EXTRACTORS_REGISTRY, f'Feature extractor "{name}" not registered')
vprint(get_kwarg('verbose', kwargs), f'Creating feature extractor "{name}" with features {list_features}')
cls = FEATURE_EXTRACTORS_REGISTRY[name]
feat_extractor = cls(name, list_features, **kwargs)
feat_extractor.eval()
if cuda:
feat_extractor.cuda()
return feat_extractor
def create_sample_similarity(name, cuda=True, **kwargs):
vassert(name in SAMPLE_SIMILARITY_REGISTRY, f'Sample similarity "{name}" not registered')
vprint(get_kwarg('verbose', kwargs), f'Creating sample similarity "{name}"')
cls = SAMPLE_SIMILARITY_REGISTRY[name]
sample_similarity = cls(name, **kwargs)
sample_similarity.eval()
if cuda:
sample_similarity.cuda()
return sample_similarity
def fid_inputs_to_metric(feat_extractor, **kwargs):
feat_layer_name = get_kwarg('feature_layer_fid', kwargs)
verbose = get_kwarg('verbose', kwargs)
vprint(verbose, f'Extracting statistics from input 1')
stats_1 = fid_input_id_to_statistics_cached(1, feat_extractor, feat_layer_name, **kwargs)
vprint(verbose, f'Extracting statistics from input 2')
stats_2 = fid_input_id_to_statistics_cached(2, feat_extractor, feat_layer_name, **kwargs)
metric = fid_statistics_to_metric(stats_1, stats_2, get_kwarg('verbose', kwargs))
return metric
def kid_features_to_metric(features_1, features_2, **kwargs):
assert torch.is_tensor(features_1) and features_1.dim() == 2
assert torch.is_tensor(features_2) and features_2.dim() == 2
assert features_1.shape[1] == features_2.shape[1]
kid_subsets = get_kwarg('kid_subsets', kwargs)
kid_subset_size = get_kwarg('kid_subset_size', kwargs)
verbose = get_kwarg('verbose', kwargs)
n_samples_1, n_samples_2 = len(features_1), len(features_2)
vassert(
n_samples_1 >= kid_subset_size and n_samples_2 >= kid_subset_size,
f'KID subset size {kid_subset_size} cannot be smaller than the number of samples (input_1: {n_samples_1}, '
f'input_2: {n_samples_2}). Consider using "kid_subset_size" kwarg or "--kid-subset-size" command line key to '
f'proceed.')
features_1 = features_1.cpu().numpy()
features_2 = features_2.cpu().numpy()
mmds = np.zeros(kid_subsets)
rng = np.random.RandomState(get_kwarg('rng_seed', kwargs))
for i in tqdm(range(kid_subsets),
disable=not verbose,
leave=False,
unit='subsets',
desc='Kernel Inception Distance'):
f1 = features_1[rng.choice(n_samples_1, kid_subset_size,
replace=False)]
f2 = features_2[rng.choice(n_samples_2, kid_subset_size,
replace=False)]
o = polynomial_mmd(
f1,
f2,
get_kwarg('kid_degree', kwargs),
get_kwarg('kid_gamma', kwargs),
get_kwarg('kid_coef0', kwargs),
)
mmds[i] = o
out = {
KEY_METRIC_KID_MEAN: float(np.mean(mmds)),
KEY_METRIC_KID_STD: float(np.std(mmds)),
}
vprint(
verbose,
f'Kernel Inception Distance: {out[KEY_METRIC_KID_MEAN]} ± {out[KEY_METRIC_KID_STD]}'
)
return out
def fid_inputs_to_metric(input_1, input_2, feat_extractor, feat_layer_name, **kwargs):
verbose = get_kwarg('verbose', kwargs)
cacheable_input1_name = get_input_cacheable_name(input_1, get_kwarg('cache_input1_name', kwargs))
cacheable_input2_name = get_input_cacheable_name(input_2, get_kwarg('cache_input2_name', kwargs))
vprint(verbose, f'Extracting statistics from input_1')
stats_1 = fid_input_to_statistics_cached(input_1, cacheable_input1_name, feat_extractor, feat_layer_name, **kwargs)
vprint(verbose, f'Extracting statistics from input_2')
stats_2 = fid_input_to_statistics_cached(input_2, cacheable_input2_name, feat_extractor, feat_layer_name, **kwargs)
metric = fid_statistics_to_metric(stats_1, stats_2, get_kwarg('verbose', kwargs))
return metric
def get_featuresdict_from_generative_model(gen_model, feat_extractor, num_samples, batch_size, cuda, rng_seed, verbose):
vassert(isinstance(gen_model, GenerativeModelBase), 'Input can only be a GenerativeModel instance')
vassert(
isinstance(feat_extractor, FeatureExtractorBase), 'Feature extractor is not a subclass of FeatureExtractorBase'
)
if batch_size > num_samples:
batch_size = num_samples
out = None
rng = np.random.RandomState(rng_seed)
if cuda:
gen_model.cuda()
with tqdm(disable=not verbose, leave=False, unit='samples', total=num_samples, desc='Processing samples') as t, \
torch.no_grad():
for sample_start in range(0, num_samples, batch_size):
sample_end = min(sample_start + batch_size, num_samples)
sz = sample_end - sample_start
noise = NOISE_SOURCE_REGISTRY[gen_model.z_type](rng, (sz, gen_model.z_size))
if cuda:
noise = noise.cuda(non_blocking=True)
gen_args = [noise]
if gen_model.num_classes > 0:
cond_labels = torch.from_numpy(rng.randint(low=0, high=gen_model.num_classes, size=(sz,), dtype=np.int))
if cuda:
cond_labels = cond_labels.cuda(non_blocking=True)
gen_args.append(cond_labels)
fakes = gen_model(*gen_args)
features = feat_extractor(fakes)
featuresdict = feat_extractor.convert_features_tuple_to_dict(features)
featuresdict = {k: [v.cpu()] for k, v in featuresdict.items()}
if out is None:
out = featuresdict
else:
out = {k: out[k] + featuresdict[k] for k in out.keys()}
t.update(sz)
vprint(verbose, 'Processing samples')
out = {k: torch.cat(v, dim=0) for k, v in out.items()}
return out
def cache_lookup_one_recompute_on_miss(cached_filename, fn_recompute, **kwargs):
if not get_kwarg('cache', kwargs):
return fn_recompute()
cache_root = get_kwarg('cache_root', kwargs)
if cache_root is None:
cache_root = os.path.join(torch.hub._get_torch_home(), 'fidelity_cache')
os.makedirs(cache_root, exist_ok=True)
item_path = os.path.join(cache_root, cached_filename + '.pt')
if os.path.exists(item_path):
vprint(get_kwarg('verbose', kwargs), f'Loading cached {item_path}')
return torch.load(item_path, map_location='cpu')
item = fn_recompute()
if get_kwarg('verbose', kwargs):
print(f'Caching {item_path}', file=sys.stderr)
torch.save(item, item_path)
return item
def isc_featuresdict_to_metric(featuresdict, feat_layer_name, **kwargs):
features = featuresdict[feat_layer_name]
out = isc_features_to_metric(
features,
get_kwarg('isc_splits', kwargs),
get_kwarg('samples_shuffle', kwargs),
get_kwarg('rng_seed', kwargs),
)
vprint(
get_kwarg('verbose', kwargs),
f'Inception Score: {out[KEY_METRIC_ISC_MEAN]} ± {out[KEY_METRIC_ISC_STD]}'
)
return out
def fid_statistics_to_metric(stat_1, stat_2, verbose):
eps = 1e-6
mu1, sigma1 = stat_1['mu'], stat_1['sigma']
mu2, sigma2 = stat_2['mu'], stat_2['sigma']
assert mu1.shape == mu2.shape and mu1.dtype == mu2.dtype
assert sigma1.shape == sigma2.shape and sigma1.dtype == sigma2.dtype
mu1 = np.atleast_1d(mu1)
mu2 = np.atleast_1d(mu2)
sigma1 = np.atleast_2d(sigma1)
sigma2 = np.atleast_2d(sigma2)
assert mu1.shape == mu2.shape, 'Training and test mean vectors have different lengths'
assert sigma1.shape == sigma2.shape, 'Training and test covariances have different dimensions'
diff = mu1 - mu2
# Product might be almost singular
covmean, _ = scipy.linalg.sqrtm(sigma1.dot(sigma2), disp=False)
if not np.isfinite(covmean).all():
vprint(verbose,
f'WARNING: fid calculation produces singular product; '
f'adding {eps} to diagonal of cov estimates'
)
offset = np.eye(sigma1.shape[0]) * eps
covmean = scipy.linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset), disp=verbose)
# Numerical error might give slight imaginary component
if np.iscomplexobj(covmean):
if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
m = np.max(np.abs(covmean.imag))
assert False, 'Imaginary component {}'.format(m)
covmean = covmean.real
tr_covmean = np.trace(covmean)
out = {
KEY_METRIC_FID: float(diff.dot(diff) + np.trace(sigma1) + np.trace(sigma2) - 2 * tr_covmean)
}
vprint(verbose, f'Frechet Inception Distance: {out[KEY_METRIC_FID]}')
return out
def get_featuresdict_from_dataset(input, feat_extractor, batch_size, cuda, save_cpu_ram, verbose):
vassert(isinstance(input, Dataset), 'Input can only be a Dataset instance')
vassert(torch.is_tensor(input[0]), 'Input Dataset should return torch.Tensor')
vassert(
isinstance(feat_extractor, FeatureExtractorBase), 'Feature extractor is not a subclass of FeatureExtractorBase'
)
if batch_size > len(input):
batch_size = len(input)
num_workers = 0 if save_cpu_ram else min(4, 2 * multiprocessing.cpu_count())
dataloader = DataLoader(
input,
batch_size=batch_size,
drop_last=False,
num_workers=num_workers,
pin_memory=cuda,
)
out = None
with tqdm(disable=not verbose, leave=False, unit='samples', total=len(input), desc='Processing samples') as t, \
torch.no_grad():
for bid, batch in enumerate(dataloader):
if cuda:
batch = batch.cuda(non_blocking=True)
features = feat_extractor(batch)
featuresdict = feat_extractor.convert_features_tuple_to_dict(features)
featuresdict = {k: [v.cpu()] for k, v in featuresdict.items()}
if out is None:
out = featuresdict
else:
out = {k: out[k] + featuresdict[k] for k in out.keys()}
t.update(batch.shape[0])
vprint(verbose, 'Processing samples')
out = {k: torch.cat(v, dim=0) for k, v in out.items()}
return out
def kid_features_to_metric(features_1, features_2, **kwargs):
verbose = get_kwarg('verbose', kwargs)
assert torch.is_tensor(features_1) and features_1.dim() == 2
assert torch.is_tensor(features_2) and features_2.dim() == 2
assert features_1.shape[1] == features_2.shape[1]
features_1 = features_1.cpu().numpy()
features_2 = features_2.cpu().numpy()
kid_subsets = get_kwarg('kid_subsets', kwargs)
kid_subset_size = get_kwarg('kid_subset_size', kwargs)
mmds = np.zeros(kid_subsets)
rng = np.random.RandomState(get_kwarg('rng_seed', kwargs))
for i in tqdm(range(kid_subsets),
disable=not verbose,
leave=False,
unit='subsets',
desc='Computing Kernel Inception Distance'):
f1 = features_1[rng.choice(len(features_1),
kid_subset_size,
replace=False)]
f2 = features_2[rng.choice(len(features_2),
kid_subset_size,
replace=False)]
o = polynomial_mmd(
f1,
f2,
get_kwarg('kid_degree', kwargs),
get_kwarg('kid_gamma', kwargs),
get_kwarg('kid_coef0', kwargs),
)
mmds[i] = o
vprint(verbose, 'Computing Kernel Inception Distance')
return {
KEY_METRIC_KID_MEAN: float(np.mean(mmds)),
KEY_METRIC_KID_STD: float(np.std(mmds)),
}
def cache_lookup_group_recompute_all_on_any_miss(cached_filename_prefix, item_names, fn_recompute, **kwargs):
verbose = get_kwarg('verbose', kwargs)
if not get_kwarg('cache', kwargs):
return fn_recompute()
cache_root = get_kwarg('cache_root', kwargs)
if cache_root is None:
cache_root = os.path.join(torch.hub._get_torch_home(), 'fidelity_cache')
os.makedirs(cache_root, exist_ok=True)
cached_paths = [os.path.join(cache_root, cached_filename_prefix + a + '.pt') for a in item_names]
if all([os.path.exists(a) for a in cached_paths]):
out = {}
for n, p in zip(item_names, cached_paths):
vprint(verbose, f'Loading cached {p}')
out[n] = torch.load(p, map_location='cpu')
return out
items = fn_recompute()
for n, p in zip(item_names, cached_paths):
vprint(verbose, f'Caching {p}')
torch.save(items[n], p)
return items
def calculate_metrics(**kwargs):
"""
Calculates metrics for the given inputs. Keyword arguments:
.. _ISC: https://arxiv.org/pdf/1606.03498.pdf
.. _FID: https://arxiv.org/pdf/1706.08500.pdf
.. _KID: https://arxiv.org/pdf/1801.01401.pdf
.. _PPL: https://arxiv.org/pdf/1812.04948.pdf
Args:
input1 (str or torch.utils.data.Dataset or GenerativeModelBase):
First input, which can be either of the following values:
- Name of a registered input. See :ref:`registry <Registry>` for the complete list of preregistered
inputs, and :meth:`register_dataset` for registering a new input. The following options refine the
behavior wrt dataset location and downloading:
:paramref:`~calculate_metrics.datasets_root`,
:paramref:`~calculate_metrics.datasets_download`.
- Path to a directory with samples. The following options refine the behavior wrt directory
traversal and samples filtering:
:paramref:`~calculate_metrics.samples_find_deep`,
:paramref:`~calculate_metrics.samples_find_ext`, and
:paramref:`~calculate_metrics.samples_ext_lossy`.
- Path to a generative model in the :obj:`ONNX<torch:torch.onnx>` or `PTH` (:obj:`JIT<torch:torch.jit>`)
format. This option also requires the following kwargs:
:paramref:`~calculate_metrics.input1_model_z_type`,
:paramref:`~calculate_metrics.input1_model_z_size`, and
:paramref:`~calculate_metrics.input1_model_num_classes`.
- Instance of :class:`~torch:torch.utils.data.Dataset` encapsulating a fixed set of samples.
- Instance of :class:`GenerativeModelBase`, implementing the generative model.
Default: `None`.
input2 (str or torch.utils.data.Dataset or GenerativeModelBase):
Second input, which can be either of the following values:
- Name of a registered input. See :ref:`registry <Registry>` for the complete list of preregistered
inputs, and :meth:`register_dataset` for registering a new input. The following options refine the
behavior wrt dataset location and downloading:
:paramref:`~calculate_metrics.datasets_root`,
:paramref:`~calculate_metrics.datasets_download`.
- Path to a directory with samples. The following options refine the behavior wrt directory
traversal and samples filtering:
:paramref:`~calculate_metrics.samples_find_deep`,
:paramref:`~calculate_metrics.samples_find_ext`, and
:paramref:`~calculate_metrics.samples_ext_lossy`.
- Path to a generative model in the :obj:`ONNX<torch:torch.onnx>` or `PTH` (:obj:`JIT<torch:torch.jit>`)
format. This option also requires the following kwargs:
:paramref:`~calculate_metrics.input2_model_z_type`,
:paramref:`~calculate_metrics.input2_model_z_size`, and
:paramref:`~calculate_metrics.input2_model_num_classes`.
- Instance of :class:`~torch:torch.utils.data.Dataset` encapsulating a fixed set of samples.
- Instance of :class:`GenerativeModelBase`, implementing the generative model.
Default: `None`.
cuda (bool): Sets executor device to GPU. Default: `True`.
batch_size (int): Batch size used to process images; the larger the more memory is used on the executor device
(see :paramref:`~calculate_metrics.cuda`). Default: `64`.
isc (bool): Calculate ISC_ (Inception Score). Default: `False`.
fid (bool): Calculate FID_ (Frechet Inception Distance). Default: `False`.
kid (bool): Calculate KID_ (Kernel Inception Distance). Default: `False`.
ppl (bool): Calculate PPL_ (Perceptual Path Length). Default: `False`.
feature_extractor (str): Name of the feature extractor (see :ref:`registry <Registry>`). Default:
`inception-v3-compat`.
feature_layer_isc (str): Name of the feature layer to use with ISC metric. Default: `logits_unbiased`.
feature_layer_fid (str): Name of the feature layer to use with FID metric. Default: `"2048"`.
feature_layer_kid (str): Name of the feature layer to use with KID metric. Default: `"2048"`.
feature_extractor_weights_path (str): Path to feature extractor weights (downloaded if `None`). Default: `None`.
isc_splits (int): Number of splits in ISC. Default: `10`.
kid_subsets (int): Number of subsets in KID. Default: `100`.
kid_subset_size (int): Subset size in KID. Default: `1000`.
kid_degree (int): Degree of polynomial kernel in KID. Default: `3`.
kid_gamma (float): Polynomial kernel gamma in KID (automatic if `None`). Default: `None`.
kid_coef0 (float): Polynomial kernel coef0 in KID. Default: `1.0`.
ppl_epsilon (float): Interpolation step size in PPL. Default: `1e-4`.
ppl_reduction (str): Reduction type to apply to the per-sample output values. Default: `mean`.
ppl_sample_similarity (str): Name of the sample similarity to use in PPL metric computation (see :ref:`registry
<Registry>`). Default: `lpips-vgg16`.
ppl_sample_similarity_resize (int): Force samples to this size when computing similarity, unless set to `None`.
Default: `64`.
ppl_sample_similarity_dtype (str): Check samples are of compatible dtype when computing similarity, unless set
to `None`. Default: `uint8`.
ppl_discard_percentile_lower (int): Removes the lower percentile of samples before reduction. Default: `1`.
ppl_discard_percentile_higher (int): Removes the higher percentile of samples before reduction. Default: `99`.
ppl_z_interp_mode (str): Noise interpolation mode in PPL (see :ref:`registry <Registry>`). Default: `lerp`.
samples_shuffle (bool): Perform random samples shuffling before computing splits. Default: `True`.
samples_find_deep (bool): Find all samples in paths recursively. Default: `False`.
samples_find_ext (str): List of comma-separated extensions (no blanks) to look for when traversing input path.
Default: `png,jpg,jpeg`.
samples_ext_lossy (str): List of comma-separated extensions (no blanks) to warn about lossy compression.
Default: `jpg,jpeg`.
datasets_root (str): Path to built-in torchvision datasets root. Default: `$ENV_TORCH_HOME/fidelity_datasets`.
datasets_download (bool): Download torchvision datasets to :paramref:`~calculate_metrics.dataset_root`.
Default: `True`.
cache_root (str): Path to file cache for features and statistics. Default: `$ENV_TORCH_HOME/fidelity_cache`.
cache (bool): Use file cache for features and statistics. Default: `True`.
input1_cache_name (str): Assigns a cache entry to input1 (when not a registered input) and forces caching of
features on it. Default: `None`.
input1_model_z_type (str): Type of noise, only required when the input is a path to a generator model (see
:ref:`registry <Registry>`). Default: `normal`.
input1_model_z_size (int): Dimensionality of noise (only required when the input is a path to a generator
model). Default: `None`.
input1_model_num_classes (int): Number of classes for conditional (0 for unconditional) generation (only
required when the input is a path to a generator model). Default: `0`.
input1_model_num_samples (int): Number of samples to draw (only required when the input is a generator model).
This option affects the following metrics: ISC, FID, KID. Default: `None`.
input2_cache_name (str): Assigns a cache entry to input2 (when not a registered input) and forces caching of
features on it. Default: `None`.
input2_model_z_type (str): Type of noise, only required when the input is a path to a generator model (see
:ref:`registry <Registry>`). Default: `normal`.
input2_model_z_size (int): Dimensionality of noise (only required when the input is a path to a generator
model). Default: `None`.
input2_model_num_classes (int): Number of classes for conditional (0 for unconditional) generation (only
required when the input is a path to a generator model). Default: `0`.
input2_model_num_samples (int): Number of samples to draw (only required when the input is a generator model).
This option affects the following metrics: FID, KID. Default: `None`.
rng_seed (int): Random numbers generator seed for all operations involving randomness. Default: `2020`.
save_cpu_ram (bool): Use less CPU RAM at the cost of speed. May not lead to improvement with every metric.
Default: `False`.
verbose (bool): Output progress information to STDERR. Default: `True`.
Returns:
: Dictionary of metrics with a subset of the following keys:
- :const:`torch_fidelity.KEY_METRIC_ISC_MEAN`
- :const:`torch_fidelity.KEY_METRIC_ISC_STD`
- :const:`torch_fidelity.KEY_METRIC_FID`
- :const:`torch_fidelity.KEY_METRIC_KID_MEAN`
- :const:`torch_fidelity.KEY_METRIC_KID_STD`
- :const:`torch_fidelity.KEY_METRIC_PPL_MEAN`
- :const:`torch_fidelity.KEY_METRIC_PPL_STD`
- :const:`torch_fidelity.KEY_METRIC_PPL_RAW`
"""
verbose = get_kwarg('verbose', kwargs)
input1, input2 = get_kwarg('input1', kwargs), get_kwarg('input2', kwargs)
have_isc = get_kwarg('isc', kwargs)
have_fid = get_kwarg('fid', kwargs)
have_kid = get_kwarg('kid', kwargs)
have_ppl = get_kwarg('ppl', kwargs)
need_input1 = have_isc or have_fid or have_kid or have_ppl
need_input2 = have_fid or have_kid
vassert(
have_isc or have_fid or have_kid or have_ppl,
'At least one of "isc", "fid", "kid", "ppl" metrics must be specified')
vassert(
input1 is not None or not need_input1,
'First input is required for "isc", "fid", "kid", and "ppl" metrics')
vassert(input2 is not None or not need_input2,
'Second input is required for "fid" and "kid" metrics')
metrics = {}
if have_isc or have_fid or have_kid:
feature_extractor = get_kwarg('feature_extractor', kwargs)
feature_layer_isc, feature_layer_fid, feature_layer_kid = (None, ) * 3
feature_layers = set()
if have_isc:
feature_layer_isc = get_kwarg('feature_layer_isc', kwargs)
feature_layers.add(feature_layer_isc)
if have_fid:
feature_layer_fid = get_kwarg('feature_layer_fid', kwargs)
feature_layers.add(feature_layer_fid)
if have_kid:
feature_layer_kid = get_kwarg('feature_layer_kid', kwargs)
feature_layers.add(feature_layer_kid)
feat_extractor = create_feature_extractor(feature_extractor,
list(feature_layers),
**kwargs)
# isc: input - featuresdict(cached) - metric
# fid: input - featuresdict(cached) - statistics(cached) - metric
# kid: input - featuresdict(cached) - metric
if (not have_isc) and have_fid and (not have_kid):
# shortcut for a case when statistics are cached and features are not required on at least one input
metric_fid = fid_inputs_to_metric(feat_extractor, **kwargs)
metrics.update(metric_fid)
else:
vprint(verbose, f'Extracting features from input1')
featuresdict_1 = extract_featuresdict_from_input_id_cached(
1, feat_extractor, **kwargs)
featuresdict_2 = None
if input2 is not None:
vprint(verbose, f'Extracting features from input2')
featuresdict_2 = extract_featuresdict_from_input_id_cached(
2, feat_extractor, **kwargs)
if have_isc:
metric_isc = isc_featuresdict_to_metric(
featuresdict_1, feature_layer_isc, **kwargs)
metrics.update(metric_isc)
if have_fid:
cacheable_input1_name = get_cacheable_input_name(1, **kwargs)
cacheable_input2_name = get_cacheable_input_name(2, **kwargs)
fid_stats_1 = fid_featuresdict_to_statistics_cached(
featuresdict_1, cacheable_input1_name, feat_extractor,
feature_layer_fid, **kwargs)
fid_stats_2 = fid_featuresdict_to_statistics_cached(
featuresdict_2, cacheable_input2_name, feat_extractor,
feature_layer_fid, **kwargs)
metric_fid = fid_statistics_to_metric(
fid_stats_1, fid_stats_2, get_kwarg('verbose', kwargs))
metrics.update(metric_fid)
if have_kid:
metric_kid = kid_featuresdict_to_metric(
featuresdict_1, featuresdict_2, feature_layer_kid,
**kwargs)
metrics.update(metric_kid)
if have_ppl:
metric_ppl = calculate_ppl(1, **kwargs)
metrics.update(metric_ppl)
return metrics
def calculate_metrics(input_1, input_2=None, **kwargs):
r"""
Calculate metrics for the given inputs.
Args:
input_1: str or torch.util.data.Dataset
First positional input, can be either a Dataset instance, or a string containing a path to a directory
of images, or one of the registered input sources (see registry.py).
input_2: str or torch.util.data.Dataset
Second positional input (not used in unary metrics, such as "isc"), can be either a Dataset instance, or a
string containing a path to a directory of images, or one of the registered input sources (see registry.py).
cuda: bool (default: True)
Sets executor device to GPU.
batch_size: int (default: 64)
Batch size used to process images; the larger the more memory is used on the executor device (see "cuda"
argument).
isc: bool (default: False)
Calculate ISC (Inception Score).
fid: bool (default: False)
Calculate FID (Frechet Inception Distance).
kid: bool (default: False)
Calculate KID (Kernel Inception Distance).
feature_extractor: str (default: inception-v3-compat)
Name of the feature extractor (see registry.py).
feature_layer_isc: str (default: logits_unbiased)
Name of the feature layer to use with ISC metric.
feature_layer_fid: str (default: 2048)
Name of the feature layer to use with FID metric.
feature_layer_kid: str (default: 2048)
Name of the feature layer to use with KID metric.
feature_extractor_weights_path: str (default: None)
Path to feature extractor weights (downloaded if None).
isc_splits: int (default: 10)
Number of splits in ISC.
kid_subsets: int (default: 100)
Number of subsets in KID.
kid_subset_size: int (default: 1000)
Subset size in KID.
kid_degree: int (default: 3)
Degree of polynomial kernel in KID.
kid_gamma: float (default: None)
Polynomial kernel gamma in KID (automatic if None).
kid_coef0: float (default: 1)
Polynomial kernel coef0 in KID.
samples_shuffle: bool (default: True)
Perform random samples shuffling before computing splits.
samples_find_deep: bool (default: False)
Find all samples in paths recursively.
samples_find_ext: str (default: png,jpg,jpeg)
List of extensions to look for when traversing input path.
samples_ext_lossy: str (default: jpg,jpeg)
List of extensions to warn about lossy compression.
datasets_root: str (default: None)
Path to built-in torchvision datasets root. Defaults to $ENV_TORCH_HOME/fidelity_datasets.
datasets_download: bool (default: True)
Download torchvision datasets to dataset_root.
cache_root: str (default: None)
Path to file cache for features and statistics. Defaults to $ENV_TORCH_HOME/fidelity_cache.
cache: bool (default: True)
Use file cache for features and statistics.
cache_input1_name: str (default: None)
Assigns a cache entry to input1 (if a path) and forces caching of features on it if not None.
cache_input2_name: str (default: None)
Assigns a cache entry to input2 (if a path) and forces caching of features on it if not None.
rng_seed: int (default: 2020)
Random numbers generator seed for all operations involving randomness.
save_cpu_ram: bool (default: False)
Use less CPU RAM at the cost of speed.
verbose: bool (default: True)
Output progress information to STDERR.
Return: a dictionary of metrics.
"""
have_isc, have_fid, have_kid = get_kwarg('isc', kwargs), get_kwarg('fid', kwargs), get_kwarg('kid', kwargs)
vassert(have_isc or have_fid or have_kid, 'At least one of "isc", "fid", "kid" metrics must be specified')
vassert(
(not have_fid) and (not have_kid) or input_2 is not None, 'Both inputs are required for "fid" and "kid" metrics'
)
verbose = get_kwarg('verbose', kwargs)
feature_layer_isc, feature_layer_fid, feature_layer_kid = (None,) * 3
feature_layers = set()
if have_isc:
feature_layer_isc = get_kwarg('feature_layer_isc', kwargs)
feature_layers.add(feature_layer_isc)
if have_fid:
feature_layer_fid = get_kwarg('feature_layer_fid', kwargs)
feature_layers.add(feature_layer_fid)
if have_kid:
feature_layer_kid = get_kwarg('feature_layer_kid', kwargs)
feature_layers.add(feature_layer_kid)
feat_extractor = create_feature_extractor(
get_kwarg('feature_extractor', kwargs), list(feature_layers), **kwargs
)
# isc: input - featuresdict(cached) - metric
# fid: input - featuresdict(cached) - statistics(cached) - metric
# kid: input - featuresdict(cached) - metric
metrics = {}
if (not have_isc) and have_fid and (not have_kid):
# shortcut for a case when statistics are cached and features are not required on at least one input
metric_fid = fid_inputs_to_metric(input_1, input_2, feat_extractor, feature_layer_fid, **kwargs)
metrics.update(metric_fid)
return metrics
cacheable_input1_name = get_input_cacheable_name(input_1, get_kwarg('cache_input1_name', kwargs))
vprint(verbose, f'Extracting features from input_1')
featuresdict_1 = extract_featuresdict_from_input_cached(input_1, cacheable_input1_name, feat_extractor, **kwargs)
featuresdict_2 = None
if input_2 is not None:
cacheable_input2_name = get_input_cacheable_name(input_2, get_kwarg('cache_input2_name', kwargs))
vprint(verbose, f'Extracting features from input_2')
featuresdict_2 = extract_featuresdict_from_input_cached(
input_2, cacheable_input2_name, feat_extractor, **kwargs
)
if have_isc:
metric_isc = isc_featuresdict_to_metric(featuresdict_1, feature_layer_isc, **kwargs)
metrics.update(metric_isc)
if have_fid:
fid_stats_1 = fid_featuresdict_to_statistics_cached(
featuresdict_1, cacheable_input1_name, feat_extractor, feature_layer_fid, **kwargs
)
fid_stats_2 = fid_featuresdict_to_statistics_cached(
featuresdict_2, cacheable_input2_name, feat_extractor, feature_layer_fid, **kwargs
)
metric_fid = fid_statistics_to_metric(fid_stats_1, fid_stats_2, get_kwarg('verbose', kwargs))
metrics.update(metric_fid)
if have_kid:
metric_kid = kid_featuresdict_to_metric(featuresdict_1, featuresdict_2, feature_layer_kid, **kwargs)
metrics.update(metric_kid)
return metrics
def calculate_ppl(input_id, **kwargs):
"""
Inspired by https://github.com/NVlabs/stylegan/blob/master/metrics/perceptual_path_length.py
"""
batch_size = get_kwarg('batch_size', kwargs)
is_cuda = get_kwarg('cuda', kwargs)
verbose = get_kwarg('verbose', kwargs)
epsilon = get_kwarg('ppl_epsilon', kwargs)
interp = get_kwarg('ppl_z_interp_mode', kwargs)
reduction = get_kwarg('ppl_reduction', kwargs)
similarity_name = get_kwarg('ppl_sample_similarity', kwargs)
sample_similarity_resize = get_kwarg('ppl_sample_similarity_resize',
kwargs)
sample_similarity_dtype = get_kwarg('ppl_sample_similarity_dtype', kwargs)
discard_percentile_lower = get_kwarg('ppl_discard_percentile_lower',
kwargs)
discard_percentile_higher = get_kwarg('ppl_discard_percentile_higher',
kwargs)
input_desc = prepare_input_descriptor_from_input_id(input_id, **kwargs)
model = prepare_input_from_descriptor(input_desc, **kwargs)
vassert(
isinstance(model, GenerativeModelBase),
'Input needs to be an instance of GenerativeModelBase, which can be either passed programmatically by wrapping '
'a model with GenerativeModelModuleWrapper, or via command line by specifying a path to a ONNX or PTH (JIT) '
'model and a set of input1_model_* arguments')
if is_cuda:
model.cuda()
input_model_num_samples = input_desc['input_model_num_samples']
input_model_num_classes = model.num_classes
input_model_z_size = model.z_size
input_model_z_type = model.z_type
vassert(
input_model_num_classes >= 0,
'Model can be unconditional (0 classes) or conditional (positive)')
vassert(
type(input_model_z_size) is int and input_model_z_size > 0,
'Dimensionality of generator noise not specified ("input1_model_z_size" argument)'
)
vassert(
type(epsilon) is float and epsilon > 0,
'Epsilon must be a small positive floating point number')
vassert(
type(input_model_num_samples) is int and input_model_num_samples > 0,
'Number of samples must be positive')
vassert(reduction in ('none', 'mean'),
'Reduction must be one of [none, mean]')
vassert(
discard_percentile_lower is None or 0 < discard_percentile_lower < 100,
'Invalid percentile')
vassert(
discard_percentile_higher is None
or 0 < discard_percentile_higher < 100, 'Invalid percentile')
if discard_percentile_lower is not None and discard_percentile_higher is not None:
vassert(0 < discard_percentile_lower < discard_percentile_higher < 100,
'Invalid percentiles')
sample_similarity = create_sample_similarity(
similarity_name,
sample_similarity_resize=sample_similarity_resize,
sample_similarity_dtype=sample_similarity_dtype,
**kwargs)
is_cond = input_desc['input_model_num_classes'] > 0
rng = np.random.RandomState(get_kwarg('rng_seed', kwargs))
lat_e0 = sample_random(rng, (input_model_num_samples, input_model_z_size),
input_model_z_type)
lat_e1 = sample_random(rng, (input_model_num_samples, input_model_z_size),
input_model_z_type)
lat_e1 = batch_interp(lat_e0, lat_e1, epsilon, interp)
labels = None
if is_cond:
labels = torch.from_numpy(
rng.randint(0, input_model_num_classes,
(input_model_num_samples, )))
distances = []
with tqdm(disable=not verbose,
leave=False,
unit='samples',
total=input_model_num_samples,
desc='Perceptual Path Length') as t, torch.no_grad():
for begin_id in range(0, input_model_num_samples, batch_size):
end_id = min(begin_id + batch_size, input_model_num_samples)
batch_sz = end_id - begin_id
batch_lat_e0 = lat_e0[begin_id:end_id]
batch_lat_e1 = lat_e1[begin_id:end_id]
if is_cond:
batch_labels = labels[begin_id:end_id]
if is_cuda:
batch_lat_e0 = batch_lat_e0.cuda(non_blocking=True)
batch_lat_e1 = batch_lat_e1.cuda(non_blocking=True)
if is_cond:
batch_labels = batch_labels.cuda(non_blocking=True)
if is_cond:
rgb_e01 = model.forward(
torch.cat((batch_lat_e0, batch_lat_e1), dim=0),
torch.cat((batch_labels, batch_labels), dim=0))
else:
rgb_e01 = model.forward(
torch.cat((batch_lat_e0, batch_lat_e1), dim=0))
rgb_e0, rgb_e1 = rgb_e01.chunk(2)
sim = sample_similarity(rgb_e0, rgb_e1)
dist_lat_e01 = sim / (epsilon**2)
distances.append(dist_lat_e01.cpu().numpy())
t.update(batch_sz)
distances = np.concatenate(distances, axis=0)
cond, lo, hi = None, None, None
if discard_percentile_lower is not None:
lo = np.percentile(distances,
discard_percentile_lower,
interpolation='lower')
cond = lo <= distances
if discard_percentile_higher is not None:
hi = np.percentile(distances,
discard_percentile_higher,
interpolation='higher')
cond = np.logical_and(cond, distances <= hi)
if cond is not None:
distances = np.extract(cond, distances)
out = {
KEY_METRIC_PPL_MEAN: float(np.mean(distances)),
KEY_METRIC_PPL_STD: float(np.std(distances))
}
if reduction == 'none':
out[KEY_METRIC_PPL_RAW] = distances
vprint(
verbose,
f'Perceptual Path Length: {out[KEY_METRIC_PPL_MEAN]} ± {out[KEY_METRIC_PPL_STD]}'
)
return out
