| import torch
|
| import numpy as np
|
| import imageio
|
|
|
| def nn_interp(H, data, sampler, shape, ema_imle, fname, logprint, preprocess_fn):
|
| mb = 10
|
| batches = []
|
| temp_latent_rnds = torch.randn([mb, H.latent_dim], dtype=torch.float32).cuda()
|
| for t in range(H.num_rows_visualize):
|
| temp_latent_rnds.normal_()
|
| tmp_snoise = [s[:mb].normal_() for s in sampler.snoise_tmp]
|
| out = ema_imle(temp_latent_rnds, tmp_snoise)
|
| batches.append((out, torch.tensor(temp_latent_rnds)))
|
| to_s = []
|
| nns = []
|
| nns_pairs = []
|
| for bb in batches:
|
| for i in range(mb):
|
| b = bb[0]
|
| to_s.append((b[i:i+1], bb[1][i:i+1]))
|
| print(len(to_s))
|
| for i in range(data.shape[0]):
|
| x = data[i:i+1]
|
| _, target = preprocess_fn([x])
|
| bst_loss = np.inf
|
| bst_ind = -1
|
| for j, dd in enumerate(to_s):
|
| d = dd[0]
|
| cur = sampler.calc_loss(target.permute(0, 3, 1, 2).cuda(), d.cuda()).item()
|
| if cur < bst_loss:
|
| bst_loss = cur
|
| bst_ind = j
|
| nns_pairs.append((bst_loss, bst_ind))
|
| nnss = torch.cat([to_s[x[1]][1] for x in nns_pairs], dim=0)
|
| torch.save(nnss.detach(), f'best-nns.npy')
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| |
