zjc
/
yolov5
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Add support for different normalization layers (#7377) 

* Add support for different normalization layers.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Cleanup

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Vardan Agarwal 2022-04-11 17:34:22 +05:30 committed by GitHub
parent bd2dda8e64
commit fa569cdae5
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
train.py
View File

@ -151,10 +151,11 @@ def train(hyp, opt, device, callbacks): # hyp is path/to/hyp.yaml or hyp dictio
LOGGER.info(f"Scaled weight_decay = {hyp['weight_decay']}") LOGGER.info(f"Scaled weight_decay = {hyp['weight_decay']}")
g = [], [], [] # optimizer parameter groups g = [], [], [] # optimizer parameter groups
bn = nn.BatchNorm2d, nn.LazyBatchNorm2d, nn.GroupNorm, nn.InstanceNorm2d, nn.LazyInstanceNorm2d, nn.LayerNorm
for v in model.modules(): for v in model.modules():
if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter): # bias if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter): # bias
g[2].append(v.bias) g[2].append(v.bias)
if isinstance(v, nn.BatchNorm2d): # weight (no decay) if isinstance(v, bn): # weight (no decay)
g[1].append(v.weight) g[1].append(v.weight)
elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter): # weight (with decay) elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter): # weight (with decay)
g[0].append(v.weight) g[0].append(v.weight)