move norm select to top

2020-08-28 08:16:07 +08:00
parent 42d6253a1d
commit 9e8e73c988
2 changed files with 17 additions and 15 deletions
--- a/model/normalization.py
+++ b/model/normalization.py
@@ -0,0 +1,13 @@
+import torch.nn as nn
+import functools
+
+
+def select_norm_layer(norm_type):
+    if norm_type == "BN":
+        return functools.partial(nn.BatchNorm2d, affine=True, track_running_stats=True)
+    elif norm_type == "IN":
+        return functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False)
+    elif norm_type == "NONE":
+        return lambda x: nn.Identity()
+    else:
+        raise NotImplemented(f'normalization layer {norm_type} is not found')