mirrors/pytorch-image-models
1
0
Fork 0
mirror of https://github.com/huggingface/pytorch-image-models.git synced 2025-06-03 15:01:08 +08:00
Code Issues Projects Releases Wiki Activity

A few minor things in SplitBN

Browse Source
This commit is contained in:
Ross Wightman 2020-01-05 20:07:03 -08:00
parent 7547119891
commit 833066b540
1 changed files with 3 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
timm/models/split_batchnorm.py
View File

@ -6,9 +6,9 @@ import torch.nn.functional as F
class SplitBatchNorm2d(torch.nn.BatchNorm2d):
def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True,
track_running_stats=True, num_splits=1):
track_running_stats=True, num_splits=2):
super().__init__(num_features, eps, momentum, affine, track_running_stats)
assert num_splits >= 2, 'Should have at least one aux BN layer (num_splits at least 2)'
assert num_splits > 1, 'Should have at least one aux BN layer (num_splits at least 2)'
self.num_splits = num_splits
self.aux_bn = nn.ModuleList([
nn.BatchNorm2d(num_features, eps, momentum, affine, track_running_stats) for _ in range(num_splits - 1)])
@ -35,8 +35,7 @@ def convert_splitbn_model(module, num_splits=2):
num_splits: number of separate batchnorm layers to split input across
Example::
>>> # model is an instance of torch.nn.Module
>>> import apex
>>> sync_bn_model = timm.models.convert_splitbn_model(model, num_splits=2)
>>> model = timm.models.convert_splitbn_model(model, num_splits=2)
"""
mod = module
if isinstance(module, torch.nn.modules.instancenorm._InstanceNorm):