ICML2022/resefa
4
1# python3.72"""Unit test for loading pre-trained models.3 4Basically, this file tests whether the perceptual model (VGG16) and the5inception model (InceptionV3), which are commonly used for loss computation and6evaluation, have the expected behavior after loading pre-trained weights. In7particular, we compare with the models from repo8 9https://github.com/NVlabs/stylegan2-ada-pytorch10"""11 12import torch13 14from models import build_model15from utils.misc import download_url16 17__all__ = ['test_model']18 19_BATCH_SIZE = 420# pylint: disable=line-too-long21_PERCEPTUAL_URL = 'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metrics/vgg16.pt'22_INCEPTION_URL = 'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metrics/inception-2015-12-05.pt'23# pylint: enable=line-too-long24 25 26def test_model():27 """Collects all model tests."""28 torch.backends.cudnn.enabled = True29 torch.backends.cudnn.allow_tf32 = False30 torch.backends.cuda.matmul.allow_tf32 = False31 torch.backends.cudnn.benchmark = False32 torch.backends.cudnn.deterministic = True33 print('========== Start Model Test ==========')34 test_perceptual()35 test_inception()36 print('========== Finish Model Test ==========')37 38 39def test_perceptual():40 """Test the perceptual model."""41 print('===== Testing Perceptual Model =====')42 43 print('Build test model.')44 model = build_model('PerceptualModel',45 use_torchvision=False,46 no_top=False,47 enable_lpips=True)48 49 print('Build reference model.')50 ref_model_path, _, = download_url(_PERCEPTUAL_URL)51 with open(ref_model_path, 'rb') as f:52 ref_model = torch.jit.load(f).eval().cuda()53 54 print('Test performance: ')55 for size in [224, 128, 256, 512, 1024]:56 raw_img = torch.randint(0, 256, size=(_BATCH_SIZE, 3, size, size))57 raw_img_comp = torch.randint(0, 256, size=(_BATCH_SIZE, 3, size, size))58 59 # The test model requires input images to have range [-1, 1].60 img = raw_img.to(torch.float32).cuda() / 127.5 - 161 img_comp = raw_img_comp.to(torch.float32).cuda() / 127.5 - 162 feat = model(img, resize_input=True, return_tensor='feature')63 pred = model(img, resize_input=True, return_tensor='prediction')64 lpips = model(img, img_comp, resize_input=False, return_tensor='lpips')65 assert feat.shape == (_BATCH_SIZE, 4096)66 assert pred.shape == (_BATCH_SIZE, 1000)67 assert lpips.shape == (_BATCH_SIZE,)68 69 # The reference model requires input images to have range [0, 255].70 img = raw_img.to(torch.float32).cuda()71 img_comp = raw_img_comp.to(torch.float32).cuda()72 ref_feat = ref_model(img, resize_images=True, return_features=True)73 ref_pred = ref_model(img, resize_images=True, return_features=False)74 temp = ref_model(torch.cat([img, img_comp], dim=0),75 resize_images=False, return_lpips=True).chunk(2)76 ref_lpips = (temp[0] - temp[1]).square().sum(dim=1, keepdim=False)77 assert ref_feat.shape == (_BATCH_SIZE, 4096)78 assert ref_pred.shape == (_BATCH_SIZE, 1000)79 assert ref_lpips.shape == (_BATCH_SIZE,)80 81 print(f' Size {size}x{size}, feature (with resize):\n '82 f'mean: {(feat - ref_feat).abs().mean().item():.3e}, '83 f'max: {(feat - ref_feat).abs().max().item():.3e}, '84 f'ref_mean: {ref_feat.abs().mean().item():.3e}, '85 f'ref_max: {ref_feat.abs().max().item():.3e}.')86 print(f' Size {size}x{size}, prediction (with resize):\n '87 f'mean: {(pred - ref_pred).abs().mean().item():.3e}, '88 f'max: {(pred - ref_pred).abs().max().item():.3e}, '89 f'ref_mean: {ref_pred.abs().mean().item():.3e}, '90 f'ref_max: {ref_pred.abs().max().item():.3e}.')91 print(f' Size {size}x{size}, LPIPS (without resize):\n '92 f'mean: {(lpips - ref_lpips).abs().mean().item():.3e}, '93 f'max: {(lpips - ref_lpips).abs().max().item():.3e}, '94 f'ref_mean: {ref_lpips.abs().mean().item():.3e}, '95 f'ref_max: {ref_lpips.abs().max().item():.3e}.')96 97 98def test_inception():99 """Test the inception model."""100 print('===== Testing Inception Model =====')101 102 print('Build test model.')103 model = build_model('InceptionModel', align_tf=True)104 105 print('Build reference model.')106 ref_model_path, _, = download_url(_INCEPTION_URL)107 with open(ref_model_path, 'rb') as f:108 ref_model = torch.jit.load(f).eval().cuda()109 110 print('Test performance: ')111 for size in [299, 128, 256, 512, 1024]:112 raw_img = torch.randint(0, 256, size=(_BATCH_SIZE, 3, size, size))113 114 # The test model requires input images to have range [-1, 1].115 img = raw_img.to(torch.float32).cuda() / 127.5 - 1116 feat = model(img)117 pred = model(img, output_predictions=True)118 pred_nb = model(img, output_predictions=True, remove_logits_bias=True)119 assert feat.shape == (_BATCH_SIZE, 2048)120 assert pred.shape == (_BATCH_SIZE, 1008)121 assert pred_nb.shape == (_BATCH_SIZE, 1008)122 123 # The reference model requires input images to have range [0, 255].124 img = raw_img.to(torch.float32).cuda()125 ref_feat = ref_model(img, return_features=True)126 ref_pred = ref_model(img)127 ref_pred_nb = ref_model(img, no_output_bias=True)128 assert ref_feat.shape == (_BATCH_SIZE, 2048)129 assert ref_pred.shape == (_BATCH_SIZE, 1008)130 assert ref_pred_nb.shape == (_BATCH_SIZE, 1008)131 132 print(f' Size {size}x{size}, feature:\n '133 f'mean: {(feat - ref_feat).abs().mean().item():.3e}, '134 f'max: {(feat - ref_feat).abs().max().item():.3e}, '135 f'ref_mean: {ref_feat.abs().mean().item():.3e}, '136 f'ref_max: {ref_feat.abs().max().item():.3e}.')137 print(f' Size {size}x{size}, prediction:\n '138 f'mean: {(pred - ref_pred).abs().mean().item():.3e}, '139 f'max: {(pred - ref_pred).abs().max().item():.3e}, '140 f'ref_mean: {ref_pred.abs().mean().item():.3e}, '141 f'ref_max: {ref_pred.abs().max().item():.3e}.')142 print(f' Size {size}x{size}, prediction (without bias):\n '143 f'mean: {(pred_nb - ref_pred_nb).abs().mean().item():.3e}, '144 f'max: {(pred_nb - ref_pred_nb).abs().max().item():.3e}, '145 f'ref_mean: {ref_pred_nb.abs().mean().item():.3e}, '146 f'ref_max: {ref_pred_nb.abs().max().item():.3e}.')147