Migrations
/
pytorch-image-models
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add drop_connect impl to try during training, fix a few comments

Browse Source
pull/6/head v0.1-weights
Ross Wightman 5 years ago
parent 0fc4cca2ff
commit 4efecfdc47
1 changed files with 38 additions and 14 deletions
Show Stats Download Patch File Download Diff File

52
models/genmobilenet.py
Unescape View File

@ -276,12 +276,13 @@ class _BlockBuilder:
"""
def __init__(self, channel_multiplier=1.0, channel_divisor=8, channel_min=None,
act_fn=None, se_gate_fn=torch.sigmoid, se_reduce_mid=False,
drop_connect_rate=0., act_fn=None, se_gate_fn=torch.sigmoid, se_reduce_mid=False,
bn_momentum=_BN_MOMENTUM_PT_DEFAULT, bn_eps=_BN_EPS_PT_DEFAULT,
folded_bn=False, padding_same=False, verbose=False):
self.channel_multiplier = channel_multiplier
self.channel_divisor = channel_divisor
self.channel_min = channel_min
self.drop_connect_rate = drop_connect_rate
self.act_fn = act_fn
self.se_gate_fn = se_gate_fn
self.se_reduce_mid = se_reduce_mid
@ -310,10 +311,12 @@ class _BlockBuilder:
print('args:', ba)
# could replace this if with lambdas or functools binding if variety increases
if bt == 'ir':
ba['drop_connect_rate'] = self.drop_connect_rate
ba['se_gate_fn'] = self.se_gate_fn
ba['se_reduce_mid'] = self.se_reduce_mid
block = InvertedResidual(**ba)
elif bt == 'ds' or bt == 'dsa':
ba['drop_connect_rate'] = self.drop_connect_rate
block = DepthwiseSeparableConv(**ba)
elif bt == 'ca':
block = CascadeConv(**ba)
@ -402,6 +405,19 @@ def hard_sigmoid(x):
return F.relu6(x + 3.) / 6.
def drop_connect(inputs, training=False, drop_connect_rate=0.):
"""Apply drop connect."""
if not training:
return inputs
keep_prob = 1 - drop_connect_rate
random_tensor = keep_prob + torch.rand(
(inputs.size()[0], 1, 1, 1), dtype=inputs.dtype, device=inputs.device)
random_tensor.floor_() # binarize
output = inputs.div(keep_prob) * random_tensor
return output
class ChannelShuffle(nn.Module):
# FIXME haven't used yet
def __init__(self, groups):
@ -474,13 +490,14 @@ class DepthwiseSeparableConv(nn.Module):
stride=1, act_fn=F.relu, noskip=False, pw_act=False,
se_ratio=0., se_gate_fn=torch.sigmoid,
bn_momentum=_BN_MOMENTUM_PT_DEFAULT, bn_eps=_BN_EPS_PT_DEFAULT,
folded_bn=False, padding_same=False):
folded_bn=False, padding_same=False, drop_connect_rate=0.):
super(DepthwiseSeparableConv, self).__init__()
assert stride in [1, 2]
self.has_se = se_ratio is not None and se_ratio > 0.
self.has_residual = (stride == 1 and in_chs == out_chs) and not noskip
self.has_pw_act = pw_act # activation after point-wise conv
self.act_fn = act_fn
self.drop_connect_rate = drop_connect_rate
dw_padding = _padding_arg(kernel_size // 2, padding_same)
pw_padding = _padding_arg(0, padding_same)
@ -515,7 +532,9 @@ class DepthwiseSeparableConv(nn.Module):
x = self.act_fn(x)
if self.has_residual:
x += residual # FIXME add drop-connect
if self.drop_connect_rate > 0.:
x = drop_connect(x, self.training, self.drop_connect_rate)
x += residual
return x
@ -557,12 +576,13 @@ class InvertedResidual(nn.Module):
se_ratio=0., se_reduce_mid=False, se_gate_fn=torch.sigmoid,
shuffle_type=None, pw_group=1,
bn_momentum=_BN_MOMENTUM_PT_DEFAULT, bn_eps=_BN_EPS_PT_DEFAULT,
folded_bn=False, padding_same=False):
folded_bn=False, padding_same=False, drop_connect_rate=0.):
super(InvertedResidual, self).__init__()
mid_chs = int(in_chs * exp_ratio)
self.has_se = se_ratio is not None and se_ratio > 0.
self.has_residual = (in_chs == out_chs and stride == 1) and not noskip
self.act_fn = act_fn
self.drop_connect_rate = drop_connect_rate
dw_padding = _padding_arg(kernel_size // 2, padding_same)
pw_padding = _padding_arg(0, padding_same)
@ -619,7 +639,9 @@ class InvertedResidual(nn.Module):
x = self.bn3(x)
if self.has_residual:
x += residual # FIXME add drop-connect
if self.drop_connect_rate > 0.:
x = drop_connect(x, self.training, self.drop_connect_rate)
x += residual
# NOTE maskrcnn_benchmark building blocks have an SE module defined here for some variants
@ -643,12 +665,14 @@ class GenMobileNet(nn.Module):
def __init__(self, block_args, num_classes=1000, in_chans=3, stem_size=32, num_features=1280,
channel_multiplier=1.0, channel_divisor=8, channel_min=None,
bn_momentum=_BN_MOMENTUM_PT_DEFAULT, bn_eps=_BN_EPS_PT_DEFAULT,
drop_rate=0., act_fn=F.relu, se_gate_fn=torch.sigmoid, se_reduce_mid=False,
drop_rate=0., drop_connect_rate=0., act_fn=F.relu,
se_gate_fn=torch.sigmoid, se_reduce_mid=False,
global_pool='avg', head_conv='default', weight_init='goog',
folded_bn=False, padding_same=False):
folded_bn=False, padding_same=False,):
super(GenMobileNet, self).__init__()
self.num_classes = num_classes
self.drop_rate = drop_rate
self.drop_connect_rate = drop_connect_rate
self.act_fn = act_fn
self.num_features = num_features
@ -661,7 +685,7 @@ class GenMobileNet(nn.Module):
builder = _BlockBuilder(
channel_multiplier, channel_divisor, channel_min,
act_fn, se_gate_fn, se_reduce_mid,
drop_connect_rate, act_fn, se_gate_fn, se_reduce_mid,
bn_momentum, bn_eps, folded_bn, padding_same, verbose=_DEBUG)
self.blocks = nn.Sequential(*builder(in_chs, block_args))
in_chs = builder.in_chs
@ -1090,7 +1114,7 @@ def _gen_spnasnet(channel_multiplier, num_classes=1000, **kwargs):
def _gen_efficientnet(channel_multiplier=1.0, depth_multiplier=1.0, num_classes=1000, **kwargs):
"""Creates a MobileNet-V3 model.
"""Creates an EfficientNet model.
Ref impl: https://github.com/tensorflow/tpu/blob/master/models/official/efficientnet/efficientnet_model.py
Paper: https://arxiv.org/abs/1905.11946
@ -1347,7 +1371,7 @@ def spnasnet_100(num_classes, in_chans=3, pretrained=False, **kwargs):
def efficientnet_b0(num_classes, in_chans=3, pretrained=False, **kwargs):
""" EfficientNet """
default_cfg = default_cfgs['efficientnet_b0']
# NOTE dropout should be 0.2 for train
# NOTE for train, drop_rate should be 0.2
model = _gen_efficientnet(
channel_multiplier=1.0, depth_multiplier=1.0,
num_classes=num_classes, in_chans=in_chans, **kwargs)
@ -1360,7 +1384,7 @@ def efficientnet_b0(num_classes, in_chans=3, pretrained=False, **kwargs):
def efficientnet_b1(num_classes, in_chans=3, pretrained=False, **kwargs):
""" EfficientNet """
default_cfg = default_cfgs['efficientnet_b1']
# NOTE dropout should be 0.2 for train
# NOTE for train, drop_rate should be 0.2
model = _gen_efficientnet(
channel_multiplier=1.0, depth_multiplier=1.1,
num_classes=num_classes, in_chans=in_chans, **kwargs)
@ -1373,7 +1397,7 @@ def efficientnet_b1(num_classes, in_chans=3, pretrained=False, **kwargs):
def efficientnet_b2(num_classes, in_chans=3, pretrained=False, **kwargs):
""" EfficientNet """
default_cfg = default_cfgs['efficientnet_b2']
# NOTE dropout should be 0.3 for train
# NOTE for train, drop_rate should be 0.3
model = _gen_efficientnet(
channel_multiplier=1.1, depth_multiplier=1.2,
num_classes=num_classes, in_chans=in_chans, **kwargs)
@ -1386,7 +1410,7 @@ def efficientnet_b2(num_classes, in_chans=3, pretrained=False, **kwargs):
def efficientnet_b3(num_classes, in_chans=3, pretrained=False, **kwargs):
""" EfficientNet """
default_cfg = default_cfgs['efficientnet_b3']
# NOTE dropout should be 0.3 for train
# NOTE for train, drop_rate should be 0.3
model = _gen_efficientnet(
channel_multiplier=1.2, depth_multiplier=1.4,
num_classes=num_classes, in_chans=in_chans, **kwargs)
@ -1399,7 +1423,7 @@ def efficientnet_b3(num_classes, in_chans=3, pretrained=False, **kwargs):
def efficientnet_b4(num_classes, in_chans=3, pretrained=False, **kwargs):
""" EfficientNet """
default_cfg = default_cfgs['efficientnet_b4']
# NOTE dropout should be 0.4 for train
# NOTE for train, drop_rate should be 0.4
model = _gen_efficientnet(
channel_multiplier=1.4, depth_multiplier=1.8,
num_classes=num_classes, in_chans=in_chans, **kwargs)

Write Preview
Loading…
Cancel
Save