(dim,
norm_layer,
in_format='channels_last',
out_format='channels_last',
eps=1e-6)
| 13 | |
| 14 | |
| 15 | def build_norm_layer(dim, |
| 16 | norm_layer, |
| 17 | in_format='channels_last', |
| 18 | out_format='channels_last', |
| 19 | eps=1e-6): |
| 20 | layers = [] |
| 21 | if norm_layer == 'BN': |
| 22 | if in_format == 'channels_last': |
| 23 | layers.append(to_channels_first()) |
| 24 | layers.append(nn.BatchNorm2d(dim)) |
| 25 | if out_format == 'channels_last': |
| 26 | layers.append(to_channels_last()) |
| 27 | elif norm_layer == 'LN': |
| 28 | if in_format == 'channels_first': |
| 29 | layers.append(to_channels_last()) |
| 30 | layers.append(nn.LayerNorm(dim, eps=eps)) |
| 31 | if out_format == 'channels_first': |
| 32 | layers.append(to_channels_first()) |
| 33 | else: |
| 34 | raise NotImplementedError( |
| 35 | f'build_norm_layer does not support {norm_layer}') |
| 36 | return nn.Sequential(*layers) |
| 37 | |
| 38 | |
| 39 | class to_channels_first(nn.Module): |
no test coverage detected