class RingLoss(nn.Module): def __init__(self, type='auto', loss_weight=1.0): """ :param type: type of loss ('l1', 'l2', 'auto') :param loss_weight: weight of loss, for 'l1' and 'l2', try with 0.01. For 'auto', try with 1.0. :return: """ super(RingLoss, self).__init__() self.radius = Parameter(torch.Tensor(1)) self.radius.data.fill_(-1) self.loss_weight = loss_weight self.type = type def forward(self, x): x = x.pow(2).sum(dim=1).pow(0.5) if self.radius.data[ 0] < 0: # Initialize the radius with the mean feature norm of first iteration self.radius.data.fill_(x.mean().data) if self.type == 'l1': # Smooth L1 Loss loss1 = F.smooth_l1_loss(x, self.radius.expand_as(x)).mul_( self.loss_weight) loss2 = F.smooth_l1_loss(self.radius.expand_as(x), x).mul_(self.loss_weight) ringloss = loss1 + loss2 elif self.type == 'auto': # Divide the L2 Loss by the feature's own norm diff = x.sub( self.radius.expand_as(x)) / (x.mean().detach().clamp(min=0.5)) diff_sq = torch.pow(torch.abs(diff), 2).mean() ringloss = diff_sq.mul_(self.loss_weight) else: # L2 Loss, if not specified diff = x.sub(self.radius.expand_as(x)) diff_sq = torch.pow(torch.abs(diff), 2).mean() ringloss = diff_sq.mul_(self.loss_weight) return ringloss
class LayerNorm(nn.Module): """ Layer Normalization based on Ba & al.: 'Layer Normalization' https://arxiv.org/pdf/1607.06450.pdf """ def __init__(self, input_size: int, learnable: bool = True, epsilon: float = 1e-6): super(LayerNorm, self).__init__() self.input_size = input_size self.learnable = learnable self.alpha = th.empty(1, input_size).fill_(0) self.beta = th.empty(1, input_size).fill_(0) self.epsilon = epsilon # Wrap as parameters if necessary if learnable: self.alpha = Parameter(self.alpha) self.beta = Parameter(self.beta) self.reset_parameters() def reset_parameters(self): std = 1.0 / math.sqrt(self.input_size) for w in self.parameters(): w.data.uniform_(-std, std) def forward(self, x: th.Tensor) -> th.Tensor: size = x.size() x = x.view(x.size(0), -1) x = (x - th.mean(x, 1).unsqueeze(1) ) / th.sqrt(th.var(x, 1).unsqueeze(1) + self.epsilon) if self.learnable: x = self.alpha.expand_as(x) * x + self.beta.expand_as(x) return x.view(size)
class BatchNorm1d(nn.Module): def __init__(self, num_features, eps=1e-5, momentum=0.1, b=True, g=True): super(BatchNorm1d, self).__init__() self.b = b self.g = g self.core = nn.BatchNorm1d(num_features, eps=eps, momentum=momentum, affine=(b and g)) print(self.core) if (not b) and g: self.g = Parameter(torch.Tensor(num_features)) elif (not g) and b: self.b = Parameter(torch.Tensor(num_features)) self.reset_parameters() def reset_parameters(self): if (not self.b) and self.g: self.g.data.fill_(1) elif (not self.g) and self.b: self.b.data.zero_() def forward(self, input): output = self.core(input) if (not self.b) and self.g: output = output * self.g.expand_as(output) elif (not self.g) and self.b: output = output + self.b.expand_as(output) return output
class CLN(nn.Module): """ Conditioned Layer Normalization """ def __init__(self, input_size, image_size, epsilon=1e-6): super(CLN, self).__init__() self.input_size = input_size self.image_size = image_size self.alpha = Tensor(1, input_size).fill_(1) self.beta = Tensor(1, input_size).fill_(0) self.epsilon = epsilon self.alpha = Parameter(self.alpha) self.beta = Parameter(self.beta) # MLP used to predict delta of alpha, beta self.fc_alpha = nn.Linear(self.image_size, self.input_size) self.fc_beta = nn.Linear(self.image_size, self.input_size) self.reset_parameters() def reset_parameters(self): std = 1.0 / math.sqrt(self.input_size) for w in self.parameters(): w.data.uniform_(-std, std) def create_cln_input(self, image_emb): delta_alpha = self.fc_alpha(image_emb) delta_beta = self.fc_beta(image_emb) return delta_alpha, delta_beta def forward(self, x, image_emb): if image_emb is None: return x # x: (batch, input_size) size = x.size() x = x.view(x.size(0), -1) x = (x - torch.mean(x, 1).unsqueeze(1).expand_as(x)) / torch.sqrt( torch.var(x, 1).unsqueeze(1).expand_as(x) + self.epsilon) delta_alpha, delta_beta = self.create_cln_input(image_emb) alpha = self.alpha.expand_as(x) + delta_alpha beta = self.beta.expand_as(x) + delta_beta x = alpha * x + beta return x.view(size)
class FullLinear(nn.Module): """ Fully connected linear readout from image-like input with c x w x h into a vector output """ def __init__(self, in_shape, outdims, bias=True): super().__init__() self.in_shape = in_shape self.outdims = outdims c, w, h = in_shape self.raw_weight = Parameter(torch.Tensor(self.outdims, c, w, h)) if bias: self.bias = Parameter(torch.Tensor(self.outdims)) else: self.register_parameter('bias', None) self.initialize() def initialize(self, init_noise=1e-3): self.raw_weight.data.normal_(0, init_noise) if self.bias is not None: self.bias.data.fill_(0) @property def weight(self): return self.raw_weight.view(self.outdims, -1) def l1(self, average=True): if average: return self.weight.abs().mean() else: return self.weight.abs().sum() def l2(self, average=True): if average: return self.weight.pow(2).mean() else: return self.weight.pow(2).sum() def forward(self, x): N = x.size(0) y = x.view(N, -1) @ self.weight.t() if self.bias is not None: y = y + self.bias.expand_as(y) return y def __repr__(self): r = self.__class__.__name__ + \ ' (' + '{} x {} x {}'.format(*self.in_shape) + ' -> ' + str(self.outdims) + ')' if self.bias is not None: r += ' with bias' return r
class ReparamNormal_MLP_Logvar(ReparamNormal): def __init__(self, input_dim=2, output_dim=2, hidden_dim=10, **kwargs): super().__init__() self.fc1 = torch.nn.Linear(input_dim, hidden_dim) self.fc2 = torch.nn.Linear(hidden_dim, hidden_dim) self.fc_mu = torch.nn.Linear(hidden_dim, output_dim, bias=False) logvar = kwargs.get('logvar', torch.zeros(output_dim)) self.logvar_param = Parameter(logvar.unsqueeze(0)) def forward(self, x): h = F.relu(self.fc1(x)) h = F.relu(self.fc2(h)) mu = self.fc_mu(h) return mu, self.logvar_param.expand_as(mu)
class Inspiration(Module): r""" Inspiration Layer (CoMatch Layer) enables the multi-style transfer in feed-forward network, which learns to match the target feature statistics during the training. This module is differentialble and can be inserted in standard feed-forward network to be learned directly from the loss function without additional supervision. .. math:: Y = \phi^{-1}[\phi(\mathcal{F}^T)W\mathcal{G}] Please see the `example of MSG-Net <./experiments/style.html>`_ training multi-style generative network for real-time transfer. Reference: Hang Zhang and Kristin Dana. "Multi-style Generative Network for Real-time Transfer." *arXiv preprint arXiv:1703.06953 (2017)* """ def __init__(self, C, B=1): super(Inspiration, self).__init__() # B is equal to 1 or input mini_batch self.weight = Parameter(torch.Tensor(1, C, C), requires_grad=True) # non-parameter buffer self.G = Variable(torch.Tensor(B, C, C), requires_grad=True) self.C = C self.reset_parameters() def reset_parameters(self): self.weight.data.uniform_(0.0, 0.02) def setTarget(self, target): self.G = target def forward(self, X): # input X is a 3D feature map self.P = torch.bmm(self.weight.expand_as(self.G), self.G) return torch.bmm( self.P.transpose(1, 2).expand(X.size(0), self.C, self.C), X.view(X.size(0), X.size(1), -1)).view_as(X) def __repr__(self): return self.__class__.__name__ + '(' \ + 'N x ' + str(self.C) + ')'
class DotProduct(torch.nn.Module): def __init__(self, in_features): super(DotProduct, self).__init__() self.in_features = in_features self.out_features = in_features self.weight = Parameter(torch.Tensor(in_features)) self.reset_parameters() def reset_parameters(self): stdv = 1. / math.sqrt(self.weight.size(0)) self.weight.data.uniform_(-stdv, stdv) #self.weight.data.normal_(0, stdv) def forward(self, input): output_np = input * self.weight.expand_as(input) return output_np def __ref__(self): return self.__class__.__name__ + '(' + 'in_features=' + str( self.in_features) + ', out_features=' + str( self.out_features) + ')'
class CoMatchLayer(Module): def __init__(self, channels, batch_size=1): super().__init__() self.C = channels self.weight = Parameter(FloatTensor(1, channels, channels), requires_grad=True) self.GM_t = FloatTensor(batch_size, channels, channels).requires_grad_() # Weight Initialization self.weight.data.uniform_(0.0, 0.02) def set_targets(self, GM_t): self.GM_t = GM_t def forward(self, x): self.P = bmm(self.weight.expand_as(self.GM_t), self.GM_t) return bmm( self.P.transpose(1, 2).expand(x.size(0), self.C, self.C), x.view(x.size(0), x.size(1), -1)).view_as(x)
class WordWindowVAE(nn.Module): def __init__(self, window_size, input_dim, embedding_dim, hidden_dim, z_dim): super(WordWindowVAE, self).__init__() self.input_dim = input_dim self.window_size = window_size self.embedding_dim = embedding_dim self.flat_dim = window_size * self.embedding_dim self.embedding = Embedding(self.input_dim, self.embedding_dim) self.logit_bias = Parameter(torch.zeros(self.input_dim)) self.fc1 = nn.Linear(self.flat_dim, hidden_dim) self.fc21 = nn.Linear(hidden_dim, z_dim) self.fc22 = nn.Linear(hidden_dim, z_dim) self.fc3 = nn.Linear(z_dim, hidden_dim) self.fc4 = nn.Linear(hidden_dim, self.flat_dim) self.relu = nn.ReLU() self.softmax = nn.Softmax(dim=-1) self.recon_loss = CrossEntropyLoss() def embed(self, x): return self.embedding(x) def encode(self, x): h1 = self.relu(self.fc1(x)) return self.fc21(h1), self.fc22(h1) def reparameterize(self, mu, logvar): if self.training: std = logvar.mul(0.5).exp_() eps = Variable(std.data.new(std.size()).normal_()) return eps.mul(std).add_(mu) else: return mu def decode(self, z): h3 = self.relu(self.fc3(z)) return self.fc4(h3) def logits(self, x): result = torch.bmm( x.view(x.size(0), -1, self.embedding_dim), self.embedding.weight.t().expand(x.size(0), self.embedding_dim, self.input_dim)) return result + self.logit_bias.expand_as(result) def forward(self, x): mu, logvar = self.encode(self.embed(x).view(-1, self.flat_dim)) z = self.reparameterize(mu, logvar) return self.decode(z), mu, logvar def loss_function(self, recon_x, x, mu, logvar): # BCE = self.recon_loss(input=self.logits(recon_x).view(-1, self.input_dim), target=x.view(-1)) target = Variable(torch.zeros(x.numel(), self.input_dim)) if torch.cuda.is_available(): target = target.cuda() target.scatter_(1, x.view(-1, 1), 1) BCE = F.binary_cross_entropy_with_logits( input=self.logits(recon_x).view(-1, self.input_dim), target=target) # see Appendix B from VAE paper: # Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014 # https://arxiv.org/abs/1312.6114 # 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2) KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp()) # Normalise by same number of elements as in reconstruction KLD /= x.numel() return BCE + KLD def get_optimizer(self, lr=1e-3): return optim.Adam(self.parameters(), lr=lr)
class BiAffineModel(nn.Module): r"""Applies a bilinear transformation to the incoming data: :math:`y = x_1 * A * x_2 + b` Args: in1_features: size of each first input sample in2_features: size of each second input sample out_features: size of each output sample bias: If set to False, the layer will not learn an additive bias. Default: ``True`` Shape: - Input: :math:`(N1, in1\_features)`, :math:`(N2, in2\_features)` - Output: :math:`(N1, N2, out\_features)` Attributes: weight: the learnable weights of the module of shape (out_features x in1_features x in2_features) bias: the learnable bias of the module of shape (out_features) Examples:: >>> m = BiAffineModel(20, 30, 40) >>> input1 = autograd.Variable(torch.randn(45, 20)) >>> input2 = autograd.Variable(torch.randn(55, 30)) >>> output = m(input1, input2) >>> print(output.size()) torch.Size([45, 55, 40]) """ def __init__(self, in1_features, in2_features, out_features, bias=True): super().__init__() self.in1_features = in1_features self.in2_features = in2_features self.out_features = out_features self.weight = Parameter( torch.Tensor(out_features, in1_features, in2_features)) if bias: self.bias = Parameter(torch.Tensor(out_features)) else: self.register_parameter('bias', None) self.reset_parameters() def reset_parameters(self): stdv = 1. / math.sqrt(self.weight.size(1)) self.weight.data.uniform_(-stdv, stdv) if self.bias is not None: self.bias.data.uniform_(-stdv, stdv) def forward(self, input1, input2): output = [] # compute output scores: for k, w in enumerate(self.weight): buff = torch.mm(input1, w) rel_score_one_role = buff.mm(input2.transpose(0, 1)) output.append( rel_score_one_role.view(rel_score_one_role.size() + (1, ))) output = torch.cat(output, -1) if self.bias is not None: output.add_(self.bias.expand_as(output)) return output.transpose(1, 2) def __repr__(self): return self.__class__.__name__ + '(' \ + 'in1_features=' + str(self.in1_features) \ + ', in2_features=' + str(self.in2_features) \ + ', out_features=' + str(self.out_features) \ + ', bias=' + str(self.bias is not None) + ')'