| 56 | |
| 57 | |
| 58 | class ScalingLayer(nn.Module): |
| 59 | def __init__(self): |
| 60 | super(ScalingLayer, self).__init__() |
| 61 | self.register_buffer("shift", torch.Tensor([-0.030, -0.088, -0.188])[None, :, None, None]) |
| 62 | self.register_buffer("scale", torch.Tensor([0.458, 0.448, 0.450])[None, :, None, None]) |
| 63 | |
| 64 | def forward(self, inp): |
| 65 | return (inp - self.shift) / self.scale |
| 66 | |
| 67 | |
| 68 | class NetLinLayer(nn.Module): |