Spaces:
Running on Zero
Running on Zero
File size: 2,384 Bytes
c7f3ffb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | import torch
from torch import nn
from torch.autograd import Function
class LayerNorm(torch.nn.LayerNorm):
"""Layer normalization module.
:param int nout: output dim size
:param int dim: dimension to be normalized
"""
def __init__(self, nout, dim=-1, eps=1e-5):
"""Construct an LayerNorm object."""
super(LayerNorm, self).__init__(nout, eps=eps)
self.dim = dim
def forward(self, x):
"""Apply layer normalization.
:param torch.Tensor x: input tensor
:return: layer normalized tensor
:rtype torch.Tensor
"""
if self.dim == -1:
return super(LayerNorm, self).forward(x)
return super(LayerNorm, self).forward(x.transpose(1, -1)).transpose(1, -1)
class Reshape(nn.Module):
def __init__(self, *args):
super(Reshape, self).__init__()
self.shape = args
def forward(self, x):
return x.view(self.shape)
class Permute(nn.Module):
def __init__(self, *args):
super(Permute, self).__init__()
self.args = args
def forward(self, x):
return x.permute(self.args)
def Linear(in_features, out_features, bias=True, init_type='xavier'):
m = nn.Linear(in_features, out_features, bias)
if init_type == 'xavier':
nn.init.xavier_uniform_(m.weight)
elif init_type == 'kaiming':
nn.init.kaiming_normal_(m.weight, mode='fan_in')
if bias:
nn.init.constant_(m.bias, 0.)
return m
def Embedding(num_embeddings, embedding_dim, padding_idx=None, init_type='normal'):
m = nn.Embedding(num_embeddings, embedding_dim, padding_idx=padding_idx)
if init_type == 'normal':
nn.init.normal_(m.weight, mean=0, std=embedding_dim ** -0.5)
elif init_type == 'kaiming':
nn.init.kaiming_normal_(m.weight, mode='fan_in')
if padding_idx is not None:
nn.init.constant_(m.weight[padding_idx], 0)
return m
class GradientReverseFunction(Function):
@staticmethod
def forward(ctx, input, coeff=1.):
ctx.coeff = coeff
output = input * 1.0
return output
@staticmethod
def backward(ctx, grad_output):
return grad_output.neg() * ctx.coeff, None
class GRL(nn.Module):
def __init__(self):
super(GRL, self).__init__()
def forward(self, *input):
return GradientReverseFunction.apply(*input)
|