|
import torch |
|
import torch.nn as nn |
|
|
|
|
|
def count_params(model): |
|
total_params = sum(p.numel() for p in model.parameters()) |
|
return total_params |
|
|
|
|
|
class ActNorm(nn.Module): |
|
def __init__(self, num_features, logdet=False, affine=True, |
|
allow_reverse_init=False): |
|
assert affine |
|
super().__init__() |
|
self.logdet = logdet |
|
self.loc = nn.Parameter(torch.zeros(1, num_features, 1, 1)) |
|
self.scale = nn.Parameter(torch.ones(1, num_features, 1, 1)) |
|
self.allow_reverse_init = allow_reverse_init |
|
|
|
self.register_buffer('initialized', torch.tensor(0, dtype=torch.uint8)) |
|
|
|
def initialize(self, input): |
|
with torch.no_grad(): |
|
flatten = input.permute(1, 0, 2, 3).contiguous().view(input.shape[1], -1) |
|
mean = ( |
|
flatten.mean(1) |
|
.unsqueeze(1) |
|
.unsqueeze(2) |
|
.unsqueeze(3) |
|
.permute(1, 0, 2, 3) |
|
) |
|
std = ( |
|
flatten.std(1) |
|
.unsqueeze(1) |
|
.unsqueeze(2) |
|
.unsqueeze(3) |
|
.permute(1, 0, 2, 3) |
|
) |
|
|
|
self.loc.data.copy_(-mean) |
|
self.scale.data.copy_(1 / (std + 1e-6)) |
|
|
|
def forward(self, input, reverse=False): |
|
if reverse: |
|
return self.reverse(input) |
|
if len(input.shape) == 2: |
|
input = input[:,:,None,None] |
|
squeeze = True |
|
else: |
|
squeeze = False |
|
|
|
_, _, height, width = input.shape |
|
|
|
if self.training and self.initialized.item() == 0: |
|
self.initialize(input) |
|
self.initialized.fill_(1) |
|
|
|
h = self.scale * (input + self.loc) |
|
|
|
if squeeze: |
|
h = h.squeeze(-1).squeeze(-1) |
|
|
|
if self.logdet: |
|
log_abs = torch.log(torch.abs(self.scale)) |
|
logdet = height*width*torch.sum(log_abs) |
|
logdet = logdet * torch.ones(input.shape[0]).to(input) |
|
return h, logdet |
|
|
|
return h |
|
|
|
def reverse(self, output): |
|
if self.training and self.initialized.item() == 0: |
|
if not self.allow_reverse_init: |
|
raise RuntimeError( |
|
"Initializing ActNorm in reverse direction is " |
|
"disabled by default. Use allow_reverse_init=True to enable." |
|
) |
|
else: |
|
self.initialize(output) |
|
self.initialized.fill_(1) |
|
|
|
if len(output.shape) == 2: |
|
output = output[:,:,None,None] |
|
squeeze = True |
|
else: |
|
squeeze = False |
|
|
|
h = output / self.scale - self.loc |
|
|
|
if squeeze: |
|
h = h.squeeze(-1).squeeze(-1) |
|
return h |
|
|
|
|
|
class AbstractEncoder(nn.Module): |
|
def __init__(self): |
|
super().__init__() |
|
|
|
def encode(self, *args, **kwargs): |
|
raise NotImplementedError |
|
|
|
|
|
class Labelator(AbstractEncoder): |
|
"""Net2Net Interface for Class-Conditional Model""" |
|
def __init__(self, n_classes, quantize_interface=True): |
|
super().__init__() |
|
self.n_classes = n_classes |
|
self.quantize_interface = quantize_interface |
|
|
|
def encode(self, c): |
|
c = c[:,None] |
|
if self.quantize_interface: |
|
return c, None, [None, None, c.long()] |
|
return c |
|
|
|
|
|
class SOSProvider(AbstractEncoder): |
|
|
|
def __init__(self, sos_token, quantize_interface=True): |
|
super().__init__() |
|
self.sos_token = sos_token |
|
self.quantize_interface = quantize_interface |
|
|
|
def encode(self, x): |
|
|
|
c = torch.ones(x.shape[0], 1)*self.sos_token |
|
c = c.long().to(x.device) |
|
if self.quantize_interface: |
|
return c, None, [None, None, c] |
|
return c |
|
|
