| 137 | |
| 138 | class TorchBatchNormLayer(nn.Module): |
| 139 | def __init__(self, n_in, params, mode, momentum=0.9, epsilon=1e-5): |
| 140 | super(TorchBatchNormLayer, self).__init__() |
| 141 | |
| 142 | scaler = params["scaler"] |
| 143 | intercept = params["intercept"] |
| 144 | |
| 145 | if mode == "1D": |
| 146 | self.layer1 = nn.BatchNorm1d( |
| 147 | num_features=n_in, momentum=1 - momentum, eps=epsilon, affine=True |
| 148 | ) |
| 149 | elif mode == "2D": |
| 150 | self.layer1 = nn.BatchNorm2d( |
| 151 | num_features=n_in, momentum=1 - momentum, eps=epsilon, affine=True |
| 152 | ) |
| 153 | |
| 154 | self.layer1.weight = nn.Parameter(torch.FloatTensor(scaler)) |
| 155 | self.layer1.bias = nn.Parameter(torch.FloatTensor(intercept)) |
| 156 | |
| 157 | def forward(self, X): |
| 158 | # (N, H, W, C) -> (N, C, H, W) |