Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| class AdaptiveInstanceNorm1d(nn.Module): | |
| def __init__(self, num_features, eps=1e-5, momentum=0.1): | |
| super(AdaptiveInstanceNorm1d, self).__init__() | |
| self.num_features = num_features | |
| self.eps = eps | |
| self.momentum = momentum | |
| self.weight = None | |
| self.bias = None | |
| self.register_buffer('running_mean', torch.zeros(num_features)) | |
| self.register_buffer('running_var', torch.ones(num_features)) | |
| def forward(self, x, direct_weighting=False, no_std=False): | |
| assert self.weight is not None and \ | |
| self.bias is not None, "Please assign AdaIN weight first" | |
| # (bs, nfeats, nframe) <= (nframe, bs, nfeats) | |
| x = x.permute(1,2,0) | |
| b, c = x.size(0), x.size(1) # batch size & channels | |
| running_mean = self.running_mean.repeat(b) | |
| running_var = self.running_var.repeat(b) | |
| # self.weight = torch.ones_like(self.weight) | |
| if direct_weighting: | |
| x_reshaped = x.contiguous().view(b * c) | |
| if no_std: | |
| out = x_reshaped + self.bias | |
| else: | |
| out = x_reshaped.mul(self.weight) + self.bias | |
| out = out.view(b, c, *x.size()[2:]) | |
| else: | |
| x_reshaped = x.contiguous().view(1, b * c, *x.size()[2:]) | |
| out = F.batch_norm( | |
| x_reshaped, running_mean, running_var, self.weight, self.bias, | |
| True, self.momentum, self.eps) | |
| out = out.view(b, c, *x.size()[2:]) | |
| # (nframe, bs, nfeats) <= (bs, nfeats, nframe) | |
| out = out.permute(2,0,1) | |
| return out | |
| def __repr__(self): | |
| return self.__class__.__name__ + '(' + str(self.num_features) + ')' | |
| def assign_adain_params(adain_params, model): | |
| # assign the adain_params to the AdaIN layers in model | |
| for m in model.modules(): | |
| if m.__class__.__name__ == "AdaptiveInstanceNorm1d": | |
| mean = adain_params[: , : m.num_features] | |
| std = adain_params[: , m.num_features: 2 * m.num_features] | |
| m.bias = mean.contiguous().view(-1) | |
| m.weight = std.contiguous().view(-1) | |
| if adain_params.size(1) > 2 * m.num_features: | |
| adain_params = adain_params[: , 2 * m.num_features:] | |
| def get_num_adain_params(model): | |
| # return the number of AdaIN parameters needed by the model | |
| num_adain_params = 0 | |
| for m in model.modules(): | |
| if m.__class__.__name__ == "AdaptiveInstanceNorm1d": | |
| num_adain_params += 2 * m.num_features | |
| return num_adain_params | |