def __init__(self,
num_images,
image_shape=None,
inception_pkl=None,
bgr2rgb=True,
inception_args=dict(normalize_input=False)):
super().__init__(num_images, image_shape=image_shape)
self.inception_pkl = inception_pkl
self.real_feats = []
self.fake_feats = []
self.real_mean = None
self.real_cov = None
self.bgr2rgb = bgr2rgb
self.device = 'cpu'
# define inception network as official StyleGAN
if inception_args.get('type', None) == 'StyleGAN':
self.inception_net = torch.jit.load(
inception_args['inception_path'])
self.inception_style = 'StyleGAN'
else:
self.inception_style = 'PyTorch'
# define inception net with default PyTorch style
self.inception_net = InceptionV3([3], **inception_args)
if torch.cuda.is_available():
self.inception_net = self.inception_net.cuda()
self.device = 'cuda'
self.inception_net.eval()
mmcv.print_log(f'FID: Adopt Inception in {self.inception_style} style',
'mmgen')
def test_fid_inception(self):
inception = InceptionV3(load_fid_inception=self.load_fid_inception)
imgs = torch.randn((2, 3, 256, 256))
out = inception(imgs)[0]
assert out.shape == (2, 2048, 1, 1)
imgs = torch.randn((2, 3, 512, 512))
out = inception(imgs)[0]
assert out.shape == (2, 2048, 1, 1)
def test_with_tero_implement(self):
self.inception = InceptionV3(load_fid_inception=True,
resize_input=False)
img = torch.randn((2, 3, 1024, 1024))
feature_ours = self.inception(img)[0].view(img.shape[0], -1)
# Tero implementation
net = torch.jit.load(
'./work_dirs/cache/inception-2015-12-05.pt').eval().cuda()
net = torch.nn.DataParallel(net)
feature_tero = net(img, return_features=True)
print(feature_ours.shape)
print((feature_tero.cpu() - feature_ours).abs().mean())
def setup_class(cls):
cls.inception = InceptionV3(load_fid_inception=False,
resize_input=True)
pipeline = [
dict(type='LoadImageFromFile', key='real_img'),
dict(
type='Resize',
keys=['real_img'],
scale=(299, 299),
keep_ratio=False,
),
dict(type='Normalize',
keys=['real_img'],
mean=[127.5] * 3,
std=[127.5] * 3,
to_rgb=False),
dict(type='Collect', keys=['real_img'], meta_keys=[]),
dict(type='ImageToTensor', keys=['real_img'])
]
dataset = UnconditionalImageDataset(
osp.join(osp.dirname(__file__), '..', 'data'), pipeline)
cls.data_loader = build_dataloader(dataset, 3, 0, dist=False)
data_loader = build_dataloader(dataset,
args.batch_size,
4,
dist=False,
shuffle=(not args.no_shuffle))
mmcv.mkdir_or_exist(args.pkl_dir)
# build inception network
if args.inception_style == 'stylegan':
inception = torch.jit.load(args.inception_pth).eval().cuda()
inception = nn.DataParallel(inception)
print_log('Adopt Inception network in StyleGAN', 'mmgen')
else:
inception = nn.DataParallel(
InceptionV3([3], resize_input=True, normalize_input=False).cuda())
inception.eval()
if args.num_samples == -1:
print_log('Use all samples in subset', 'mmgen')
num_samples = len(dataset)
else:
num_samples = args.num_samples
features = extract_inception_features(data_loader, inception, num_samples,
args.inception_style).numpy()
# sanity check for the number of features
assert features.shape[
0] == num_samples, 'the number of features != num_samples'
print_log(f'Extract {num_samples} features', 'mmgen')