From 996c77aa9405a2481e54a3ce5ca7def73d3e9134 Mon Sep 17 00:00:00 2001 From: Ross Wightman Date: Mon, 15 Apr 2019 16:58:40 -0700 Subject: [PATCH] Prep mnasnet for pretrained models, use the select global pool, some comment mistakes --- models/mnasnet.py | 43 +++++++++++++++++++++++++++++++------------ 1 file changed, 31 insertions(+), 12 deletions(-) diff --git a/models/mnasnet.py b/models/mnasnet.py index 00887733..1a2a2990 100644 --- a/models/mnasnet.py +++ b/models/mnasnet.py @@ -29,7 +29,7 @@ def _cfg(url='', **kwargs): 'url': url, 'num_classes': 1000, 'input_size': (3, 224, 224), 'pool_size': (7, 7), 'crop_pct': 0.875, 'interpolation': 'bilinear', 'mean': IMAGENET_DEFAULT_MEAN, 'std': IMAGENET_DEFAULT_STD, - 'first_conv': 'layer0.conv1', 'classifier': 'last_linear', + 'first_conv': 'conv_stem', 'classifier': 'classifier', **kwargs } @@ -207,12 +207,11 @@ class MnasNet(nn.Module): global_pool='avg', act_fn=F.relu): super(MnasNet, self).__init__() self.num_classes = num_classes - self.drop_rate = drop_rate - self.global_pool = global_pool - self.act_fn = act_fn self.depth_multiplier = depth_multiplier self.bn_momentum = bn_momentum self.bn_eps = bn_eps + self.drop_rate = drop_rate + self.act_fn = act_fn self.num_features = 1280 self.conv_stem = nn.Conv2d(in_chans, stem_size, 3, padding=1, stride=2, bias=False) @@ -230,7 +229,7 @@ class MnasNet(nn.Module): self.conv_head = nn.Conv2d(out_chs, self.num_features, 1, padding=0, stride=1, bias=False) self.bn1 = nn.BatchNorm2d(self.num_features, momentum=self.bn_momentum, eps=self.bn_eps) - self.avg_pool = nn.AdaptiveAvgPool2d(1) + self.global_pool = SelectAdaptivePool2d(pool_type=global_pool) self.classifier = nn.Linear(self.num_features, self.num_classes) for m in self.modules(): @@ -251,11 +250,12 @@ class MnasNet(nn.Module): return self.classifier def reset_classifier(self, num_classes, global_pool='avg'): - #self.global_pool = SelectAdaptivePool2d(pool_type=global_pool) + self.global_pool = SelectAdaptivePool2d(pool_type=global_pool) self.num_classes = num_classes del self.classifier if num_classes: - self.classifier = nn.Linear(self.num_features, num_classes) + self.classifier = nn.Linear( + self.num_features * self.global_pool.feat_mult(), num_classes) else: self.classifier = None @@ -282,7 +282,7 @@ class MnasNet(nn.Module): def mnasnet_a1(depth_multiplier, num_classes=1000, **kwargs): """Creates a mnasnet-a1 model. Args: - depth_multiplier: multiplier to number of filters per layer. + depth_multiplier: multiplier to number of channels per layer. """ # defs from https://github.com/tensorflow/tpu/blob/master/models/official/mnasnet/mnasnet_models.py @@ -314,7 +314,7 @@ def mnasnet_a1(depth_multiplier, num_classes=1000, **kwargs): def mnasnet_b1(depth_multiplier, num_classes=1000, **kwargs): """Creates a mnasnet-b1 model. Args: - depth_multiplier: multiplier to number of filters per layer. + depth_multiplier: multiplier to number of channels per layer. """ # from https://github.com/tensorflow/tpu/blob/master/models/official/mnasnet/mnasnet_models.py blocks_defs = [ @@ -345,7 +345,7 @@ def mnasnet_b1(depth_multiplier, num_classes=1000, **kwargs): def mnasnet_small(depth_multiplier, num_classes=1000, **kwargs): """Creates a mnasnet-b1 model. Args: - depth_multiplier: multiplier to number of filters per layer. + depth_multiplier: multiplier to number of channels per layer. """ # from https://github.com/tensorflow/tpu/blob/master/models/official/mnasnet/mnasnet_models.py blocks_defs = [ @@ -378,6 +378,8 @@ def mnasnet0_50(num_classes=1000, in_chans=3, pretrained=False, **kwargs): default_cfg = default_cfgs['mnasnet0_50'] model = mnasnet_b1(0.5, num_classes=num_classes, in_chans=in_chans, **kwargs) model.default_cfg = default_cfg + if pretrained: + load_pretrained(model, default_cfg, num_classes, in_chans) return model @@ -386,6 +388,8 @@ def mnasnet0_75(num_classes, in_chans=3, pretrained=False, **kwargs): default_cfg = default_cfgs['mnasnet0_50'] model = mnasnet_b1(0.75, num_classes=num_classes, in_chans=in_chans, **kwargs) model.default_cfg = default_cfg + if pretrained: + load_pretrained(model, default_cfg, num_classes, in_chans) return model @@ -394,6 +398,8 @@ def mnasnet1_00(num_classes, in_chans=3, pretrained=False, **kwargs): default_cfg = default_cfgs['mnasnet0_50'] model = mnasnet_b1(1.0, num_classes=num_classes, in_chans=in_chans, **kwargs) model.default_cfg = default_cfg + if pretrained: + load_pretrained(model, default_cfg, num_classes, in_chans) return model @@ -402,6 +408,8 @@ def mnasnet1_40(num_classes, in_chans=3, pretrained=False, **kwargs): default_cfg = default_cfgs['mnasnet0_50'] model = mnasnet_b1(1.4, num_classes=num_classes, in_chans=in_chans, **kwargs) model.default_cfg = default_cfg + if pretrained: + load_pretrained(model, default_cfg, num_classes, in_chans) return model @@ -410,6 +418,8 @@ def semnasnet0_50(num_classes=1000, in_chans=3, pretrained=False, **kwargs): default_cfg = default_cfgs['mnasnet0_50'] model = mnasnet_a1(0.5, num_classes=num_classes, in_chans=in_chans, **kwargs) model.default_cfg = default_cfg + if pretrained: + load_pretrained(model, default_cfg, num_classes, in_chans) return model @@ -418,27 +428,36 @@ def semnasnet0_75(num_classes, in_chans=3, pretrained=False, **kwargs): default_cfg = default_cfgs['mnasnet0_50'] model = mnasnet_a1(0.75, num_classes=num_classes, in_chans=in_chans, **kwargs) model.default_cfg = default_cfg + if pretrained: + load_pretrained(model, default_cfg, num_classes, in_chans) return model def semnasnet1_00(num_classes, in_chans=3, pretrained=False, **kwargs): - """ MNASNet Small, depth multiplier of 1.0. """ + """ MNASNet A1 (w/ SE), depth multiplier of 1.0. """ default_cfg = default_cfgs['mnasnet0_50'] model = mnasnet_a1(1.0, num_classes=num_classes, in_chans=in_chans, **kwargs) model.default_cfg = default_cfg + if pretrained: + load_pretrained(model, default_cfg, num_classes, in_chans) return model def semnasnet1_40(num_classes, in_chans=3, pretrained=False, **kwargs): - """ MNASNet with depth multiplier of 1.3. """ + """ MNASNet A1 (w/ SE), depth multiplier of 1.4. """ default_cfg = default_cfgs['mnasnet0_50'] model = mnasnet_a1(1.4, num_classes=num_classes, in_chans=in_chans, **kwargs) model.default_cfg = default_cfg + if pretrained: + load_pretrained(model, default_cfg, num_classes, in_chans) return model def mnasnet_small(num_classes, in_chans=3, pretrained=False, **kwargs): + """ MNASNet Small, depth multiplier of 1.0. """ default_cfg = default_cfgs['mnasnet_small'] model = mnasnet_small(1.0, num_classes=num_classes, in_chans=in_chans, **kwargs) model.default_cfg = default_cfg + if pretrained: + load_pretrained(model, default_cfg, num_classes, in_chans) return model