| 84 | |
| 85 | |
| 86 | class InstantLayerNorm2d(nn.Module): |
| 87 | def __init__(self, |
| 88 | num_features, |
| 89 | affine=True, |
| 90 | eps=1e-5, |
| 91 | ): |
| 92 | super(InstantLayerNorm2d, self).__init__() |
| 93 | self.num_features = num_features |
| 94 | self.affine = affine |
| 95 | self.eps = eps |
| 96 | if affine: |
| 97 | self.gain = nn.Parameter(torch.ones(1, num_features, 1, 1), requires_grad=True) |
| 98 | self.bias = nn.Parameter(torch.zeros(1, num_features, 1, 1), requires_grad=True) |
| 99 | else: |
| 100 | self.gain = Variable(torch.ones(1, num_features, 1, 1), requires_grad=False) |
| 101 | self.bias = Variable(torch.zeros(1, num_features, 1, 1), requires_grad=False) |
| 102 | |
| 103 | def forward(self, inpt): |
| 104 | # inpt: (B,C,T,F) |
| 105 | ins_mean = torch.mean(inpt, dim=[1,3], keepdim=True) # (B,C,T,1) |
| 106 | ins_std = (torch.std(inpt, dim=[1,3], keepdim=True) + self.eps).pow(0.5) # (B,C,T,1) |
| 107 | x = (inpt - ins_mean) / ins_std |
| 108 | return x * self.gain.expand_as(x).type(x.type()) + self.bias.expand_as(x).type(x.type()) |
| 109 | |
| 110 | |
| 111 | |