Combine Instance Norm and Batch Norm into One Layer
| 18 | #------------------------------------------------------------------------------ |
| 19 | |
| 20 | class IBNorm(nn.Module): |
| 21 | """ Combine Instance Norm and Batch Norm into One Layer |
| 22 | """ |
| 23 | |
| 24 | def __init__(self, in_channels): |
| 25 | super(IBNorm, self).__init__() |
| 26 | in_channels = in_channels |
| 27 | self.bnorm_channels = int(in_channels / 2) |
| 28 | self.inorm_channels = in_channels - self.bnorm_channels |
| 29 | |
| 30 | self.bnorm = nn.BatchNorm2d(self.bnorm_channels, affine=True) |
| 31 | self.inorm = nn.InstanceNorm2d(self.inorm_channels, affine=False) |
| 32 | |
| 33 | def forward(self, x): |
| 34 | bn_x = self.bnorm(x[:, :self.bnorm_channels, ...].contiguous()) |
| 35 | in_x = self.inorm(x[:, self.bnorm_channels:, ...].contiguous()) |
| 36 | |
| 37 | return torch.cat((bn_x, in_x), 1) |
| 38 | |
| 39 | |
| 40 | class Conv2dIBNormRelu(nn.Module): |