glenn-jocher commited on
Commit
4bb7eb8
·
unverified ·
1 Parent(s): fa569cd

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

Files changed (1) hide show
  1. train.py +1 -1
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 = nn.BatchNorm2d, nn.LazyBatchNorm2d, nn.GroupNorm, nn.InstanceNorm2d, nn.LazyInstanceNorm2d, nn.LayerNorm
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)