Spaces:
Runtime error
Runtime error
File size: 2,992 Bytes
97069e1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 |
import argparse
import torch
from torch.nn import functional as F
import numpy as np
from tqdm import tqdm
import lpips
from model import Generator
def normalize(x):
return x / torch.sqrt(x.pow(2).sum(-1, keepdim=True))
def slerp(a, b, t):
a = normalize(a)
b = normalize(b)
d = (a * b).sum(-1, keepdim=True)
p = t * torch.acos(d)
c = normalize(b - d * a)
d = a * torch.cos(p) + c * torch.sin(p)
return normalize(d)
def lerp(a, b, t):
return a + (b - a) * t
if __name__ == '__main__':
device = 'cuda'
parser = argparse.ArgumentParser()
parser.add_argument('--space', choices=['z', 'w'])
parser.add_argument('--batch', type=int, default=64)
parser.add_argument('--n_sample', type=int, default=5000)
parser.add_argument('--size', type=int, default=256)
parser.add_argument('--eps', type=float, default=1e-4)
parser.add_argument('--crop', action='store_true')
parser.add_argument('ckpt', metavar='CHECKPOINT')
args = parser.parse_args()
latent_dim = 512
ckpt = torch.load(args.ckpt)
g = Generator(args.size, latent_dim, 8).to(device)
g.load_state_dict(ckpt['g_ema'])
g.eval()
percept = lpips.PerceptualLoss(
model='net-lin', net='vgg', use_gpu=device.startswith('cuda')
)
distances = []
n_batch = args.n_sample // args.batch
resid = args.n_sample - (n_batch * args.batch)
batch_sizes = [args.batch] * n_batch + [resid]
with torch.no_grad():
for batch in tqdm(batch_sizes):
noise = g.make_noise()
inputs = torch.randn([batch * 2, latent_dim], device=device)
lerp_t = torch.rand(batch, device=device)
if args.space == 'w':
latent = g.get_latent(inputs)
latent_t0, latent_t1 = latent[::2], latent[1::2]
latent_e0 = lerp(latent_t0, latent_t1, lerp_t[:, None])
latent_e1 = lerp(latent_t0, latent_t1, lerp_t[:, None] + args.eps)
latent_e = torch.stack([latent_e0, latent_e1], 1).view(*latent.shape)
image, _ = g([latent_e], input_is_latent=True, noise=noise)
if args.crop:
c = image.shape[2] // 8
image = image[:, :, c * 3 : c * 7, c * 2 : c * 6]
factor = image.shape[2] // 256
if factor > 1:
image = F.interpolate(
image, size=(256, 256), mode='bilinear', align_corners=False
)
dist = percept(image[::2], image[1::2]).view(image.shape[0] // 2) / (
args.eps ** 2
)
distances.append(dist.to('cpu').numpy())
distances = np.concatenate(distances, 0)
lo = np.percentile(distances, 1, interpolation='lower')
hi = np.percentile(distances, 99, interpolation='higher')
filtered_dist = np.extract(
np.logical_and(lo <= distances, distances <= hi), distances
)
print('ppl:', filtered_dist.mean())
|