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) + ')'
