| 57 | raise ValueError('No checkpoint found at "%s"' % filename) |
| 58 | |
| 59 | class InstantLayerNorm1d(nn.Module): |
| 60 | def __init__(self, |
| 61 | num_features, |
| 62 | affine=True, |
| 63 | eps=1e-5, |
| 64 | ): |
| 65 | super(InstantLayerNorm1d, self).__init__() |
| 66 | self.num_features = num_features |
| 67 | self.affine = affine |
| 68 | self.eps = eps |
| 69 | |
| 70 | if affine: |
| 71 | self.gain = nn.Parameter(torch.ones(1, 1, num_features), requires_grad=True) |
| 72 | self.bias = nn.Parameter(torch.zeros(1, 1, num_features), requires_grad=True) |
| 73 | else: |
| 74 | self.gain = Variable(torch.ones(1, 1, num_features), requires_grad=False) |
| 75 | self.bias = Variable(torch.zeros(1, 1, num_features), requires_gra=False) |
| 76 | |
| 77 | def forward(self, inpt): |
| 78 | # inpt: (T,B,C) |
| 79 | seq_len, b_size, channel = inpt.shape |
| 80 | ins_mean = torch.mean(inpt, dim=-1, keepdim=True) # (T,B,1) |
| 81 | ins_std = (torch.var(inpt, dim=-1, keepdim=True) + self.eps).pow(0.5) # (T,B,1) |
| 82 | x = (inpt - ins_mean) / ins_std |
| 83 | return x * self.gain.expand_as(x).type(x.type()) + self.bias.expand_as(x).type(x.type()) |
| 84 | |
| 85 | |
| 86 | class InstantLayerNorm2d(nn.Module): |