Dynamic normalization layer selection (#7392)
Browse files* Dynamic normalization layer selection
Based on actual available layers. Torch 1.7 compatible, resolves https://github.com/ultralytics/yolov5/issues/7381
* Update train.py
train.py
CHANGED
@@ -151,7 +151,7 @@ def train(hyp, opt, device, callbacks): # hyp is path/to/hyp.yaml or hyp dictio
|
|
151 |
LOGGER.info(f"Scaled weight_decay = {hyp['weight_decay']}")
|
152 |
|
153 |
g = [], [], [] # optimizer parameter groups
|
154 |
-
bn =
|
155 |
for v in model.modules():
|
156 |
if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter): # bias
|
157 |
g[2].append(v.bias)
|
|
|
151 |
LOGGER.info(f"Scaled weight_decay = {hyp['weight_decay']}")
|
152 |
|
153 |
g = [], [], [] # optimizer parameter groups
|
154 |
+
bn = tuple(v for k, v in nn.__dict__.items() if 'Norm' in k) # normalization layers, i.e. BatchNorm2d()
|
155 |
for v in model.modules():
|
156 |
if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter): # bias
|
157 |
g[2].append(v.bias)
|