Spaces:
Runtime error
Runtime error
| import torch | |
| import torch.nn as nn | |
| from models.util import mean_dim | |
| class _BaseNorm(nn.Module): | |
| """Base class for ActNorm (Glow) and PixNorm (Flow++). | |
| The mean and inv_std get initialized using the mean and variance of the | |
| first mini-batch. After the init, mean and inv_std are trainable parameters. | |
| Adapted from: | |
| > https://github.com/openai/glow | |
| """ | |
| def __init__(self, num_channels, height, width): | |
| super(_BaseNorm, self).__init__() | |
| # Input gets concatenated along channel axis | |
| #num_channels *= 2 | |
| self.register_buffer('is_initialized', torch.zeros(1)) | |
| self.mean = nn.Parameter(torch.zeros(1, num_channels, height, width)) | |
| self.inv_std = nn.Parameter(torch.zeros(1, num_channels, height, width)) | |
| self.eps = 1e-6 | |
| def initialize_parameters(self, x): | |
| if not self.training: | |
| return | |
| with torch.no_grad(): | |
| mean, inv_std = self._get_moments(x) | |
| self.mean.data.copy_(mean.data) | |
| self.inv_std.data.copy_(inv_std.data) | |
| self.is_initialized += 1. | |
| def _center(self, x, reverse=False): | |
| if reverse: | |
| return x + self.mean | |
| else: | |
| return x - self.mean | |
| def _get_moments(self, x): | |
| raise NotImplementedError('Subclass of _BaseNorm must implement _get_moments') | |
| def _scale(self, x, sldj, reverse=False): | |
| raise NotImplementedError('Subclass of _BaseNorm must implement _scale') | |
| def forward(self, x, cond, ldj=None, reverse=False): | |
| #import pdb;pdb.set_trace() | |
| x = torch.cat(x, dim=1) | |
| # import pdb;pdb.set_trace() | |
| if not self.is_initialized: | |
| print("Initializing norm Layer!") | |
| self.initialize_parameters(x) | |
| if reverse: | |
| x, ldj = self._scale(x, ldj, reverse) | |
| x = self._center(x, reverse) | |
| else: | |
| x = self._center(x, reverse) | |
| x, ldj = self._scale(x, ldj, reverse) | |
| x = x.chunk(2, dim=1) | |
| return x, ldj | |
| class BatchNorm(nn.Module): | |
| def __init__(self, num_channels, momentum=0.1): | |
| super(BatchNorm, self).__init__() | |
| self.gamma = nn.Parameter(torch.ones(1, num_channels, 1, 1)) | |
| self.beta = nn.Parameter(torch.zeros(1, num_channels, 1, 1)) | |
| self.register_buffer('running_mean', torch.zeros(1, num_channels, 1, 1)) | |
| self.register_buffer('running_var', torch.ones(1, num_channels, 1, 1)) | |
| self.eps = 1e-5 | |
| self.momentum = momentum | |
| self.inv_std = None | |
| self.register_buffer('is_initialized', torch.zeros(1)) | |
| def _get_moments(self, x): | |
| mean = mean_dim(x.clone(), dim=[0, 2, 3], keepdims=True).detach() | |
| var = mean_dim((x.clone() - mean) ** 2, dim=[0, 2, 3], keepdims=True).detach() | |
| # inv_std = 1. / (var.sqrt() + self.eps) | |
| if not self.is_initialized: | |
| self.running_mean.data.copy_(mean.data) | |
| self.running_var.data.copy_(var.data) | |
| self.is_initialized += 1. | |
| else: | |
| if self.momentum < 1.0: | |
| self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * mean | |
| self.running_var = (1 - self.momentum) * self.running_var + self.momentum * var | |
| else: | |
| self.running_mean.data.copy_(mean.data) | |
| self.running_var.data.copy_(var.data) | |
| def forward(self, x, cond, ldj=None, reverse=False): | |
| # import pdb;pdb.set_trace() | |
| x = torch.cat(x, dim=1) | |
| if self.training: | |
| # print("HI") | |
| self._get_moments(x) | |
| # print(self.running_var[0]) | |
| inv_std = 1. / (self.running_var.sqrt() + self.eps) | |
| if reverse: | |
| x = self._center(x, self.beta, reverse) | |
| x, ldj = self._scale(x, ldj, self.gamma, reverse) | |
| x, ldj = self._scale(x, ldj, inv_std, reverse) | |
| x = self._center(x, self.running_mean, reverse) | |
| else: | |
| x = self._center(x, self.running_mean, reverse) | |
| x, ldj = self._scale(x, ldj, inv_std, reverse) | |
| x, ldj = self._scale(x, ldj, self.gamma, reverse) | |
| x = self._center(x, self.beta, reverse) | |
| x = x.chunk(2, dim=1) | |
| return x, ldj | |
| def _center(self, x, centerer, reverse=False): | |
| if reverse: | |
| return x + centerer | |
| else: | |
| return x - centerer | |
| def _scale(self, x, sldj, scaler, reverse=False): | |
| if reverse: | |
| x = x / scaler | |
| sldj = sldj - scaler.log().sum() * x.size(2) * x.size(3) | |
| else: | |
| x = x * scaler | |
| sldj = sldj + scaler.log().sum() * x.size(2) * x.size(3) | |
| return x, sldj | |
| class ActNorm(_BaseNorm): | |
| """Activation Normalization used in Glow | |
| The mean and inv_std get initialized using the mean and variance of the | |
| first mini-batch. After the init, mean and inv_std are trainable parameters. | |
| """ | |
| def __init__(self, num_channels): | |
| super(ActNorm, self).__init__(num_channels, 1, 1) | |
| def _get_moments(self, x): | |
| mean = mean_dim(x.clone(), dim=[0, 2, 3], keepdims=True) | |
| var = mean_dim((x.clone() - mean) ** 2, dim=[0, 2, 3], keepdims=True) | |
| inv_std = 1. / (var.sqrt() + self.eps) | |
| return mean, inv_std | |
| def _scale(self, x, sldj, reverse=False): | |
| if reverse: | |
| x = x / self.inv_std | |
| sldj = sldj - self.inv_std.log().sum() * x.size(2) * x.size(3) | |
| else: | |
| x = x * self.inv_std | |
| sldj = sldj + self.inv_std.log().sum() * x.size(2) * x.size(3) | |
| return x, sldj | |
| class PixNorm(_BaseNorm): | |
| """Pixel-wise Activation Normalization used in Flow++ | |
| Normalizes every activation independently (note this differs from the variant | |
| used in in Glow, where they normalize each channel). The mean and stddev get | |
| initialized using the mean and stddev of the first mini-batch. After the | |
| initialization, `mean` and `inv_std` become trainable parameters. | |
| """ | |
| def _get_moments(self, x): | |
| mean = torch.mean(x.clone(), dim=0, keepdim=True) | |
| var = torch.mean((x.clone() - mean) ** 2, dim=0, keepdim=True) | |
| inv_std = 1. / (var.sqrt() + self.eps) | |
| return mean, inv_std | |
| def _scale(self, x, sldj, reverse=False): | |
| if reverse: | |
| x = x / self.inv_std | |
| sldj = sldj - self.inv_std.log().sum() | |
| else: | |
| x = x * self.inv_std | |
| sldj = sldj + self.inv_std.log().sum() | |
| return x, sldj | |