Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,704 Bytes
d65c9b3 |
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 |
import torch
from random import randrange
import torch.nn.functional as F
def noise_regularization(
e_t, noise_pred_optimal, lambda_kl, lambda_ac, num_reg_steps, num_ac_rolls
):
for _outer in range(num_reg_steps):
if lambda_kl > 0:
_var = torch.autograd.Variable(e_t.detach().clone(), requires_grad=True)
l_kld = patchify_latents_kl_divergence(_var, noise_pred_optimal)
l_kld.backward()
_grad = _var.grad.detach()
_grad = torch.clip(_grad, -100, 100)
e_t = e_t - lambda_kl * _grad
if lambda_ac > 0:
for _inner in range(num_ac_rolls):
_var = torch.autograd.Variable(e_t.detach().clone(), requires_grad=True)
l_ac = auto_corr_loss(_var)
l_ac.backward()
_grad = _var.grad.detach() / num_ac_rolls
e_t = e_t - lambda_ac * _grad
e_t = e_t.detach()
return e_t
def auto_corr_loss(x, random_shift=True):
B, C, H, W = x.shape
assert B == 1
x = x.squeeze(0)
# x must be shape [C,H,W] now
reg_loss = 0.0
for ch_idx in range(x.shape[0]):
noise = x[ch_idx][None, None, :, :]
while True:
if random_shift:
roll_amount = randrange(noise.shape[2] // 2)
else:
roll_amount = 1
reg_loss += (
noise * torch.roll(noise, shifts=roll_amount, dims=2)
).mean() ** 2
reg_loss += (
noise * torch.roll(noise, shifts=roll_amount, dims=3)
).mean() ** 2
if noise.shape[2] <= 8:
break
noise = F.avg_pool2d(noise, kernel_size=2)
return reg_loss
def patchify_latents_kl_divergence(x0, x1, patch_size=4, num_channels=4):
def patchify_tensor(input_tensor):
patches = (
input_tensor.unfold(1, patch_size, patch_size)
.unfold(2, patch_size, patch_size)
.unfold(3, patch_size, patch_size)
)
patches = patches.contiguous().view(-1, num_channels, patch_size, patch_size)
return patches
x0 = patchify_tensor(x0)
x1 = patchify_tensor(x1)
kl = latents_kl_divergence(x0, x1).sum()
return kl
def latents_kl_divergence(x0, x1):
EPSILON = 1e-6
x0 = x0.view(x0.shape[0], x0.shape[1], -1)
x1 = x1.view(x1.shape[0], x1.shape[1], -1)
mu0 = x0.mean(dim=-1)
mu1 = x1.mean(dim=-1)
var0 = x0.var(dim=-1)
var1 = x1.var(dim=-1)
kl = (
torch.log((var1 + EPSILON) / (var0 + EPSILON))
+ (var0 + (mu0 - mu1) ** 2) / (var1 + EPSILON)
- 1
)
kl = torch.abs(kl).sum(dim=-1)
return kl
|