release code
This commit is contained in:
miunangel
217
Dassl.ProGrad.pytorch/dassl/modeling/backbone/mobilenetv2.py
Normal file
217
Dassl.ProGrad.pytorch/dassl/modeling/backbone/mobilenetv2.py
Normal file
@@ -0,0 +1,217 @@
|
||||
import torch.utils.model_zoo as model_zoo
|
||||
from torch import nn
|
||||
|
||||
from .build import BACKBONE_REGISTRY
|
||||
from .backbone import Backbone
|
||||
|
||||
model_urls = {
|
||||
"mobilenet_v2":
|
||||
"https://download.pytorch.org/models/mobilenet_v2-b0353104.pth",
|
||||
}
|
||||
|
||||
|
||||
def _make_divisible(v, divisor, min_value=None):
|
||||
"""
|
||||
This function is taken from the original tf repo.
|
||||
It ensures that all layers have a channel number that is divisible by 8
|
||||
It can be seen here:
|
||||
https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py
|
||||
:param v:
|
||||
:param divisor:
|
||||
:param min_value:
|
||||
:return:
|
||||
"""
|
||||
if min_value is None:
|
||||
min_value = divisor
|
||||
new_v = max(min_value, int(v + divisor/2) // divisor * divisor)
|
||||
# Make sure that round down does not go down by more than 10%.
|
||||
if new_v < 0.9 * v:
|
||||
new_v += divisor
|
||||
return new_v
|
||||
|
||||
|
||||
class ConvBNReLU(nn.Sequential):
|
||||
|
||||
def __init__(
|
||||
self, in_planes, out_planes, kernel_size=3, stride=1, groups=1
|
||||
):
|
||||
padding = (kernel_size-1) // 2
|
||||
super().__init__(
|
||||
nn.Conv2d(
|
||||
in_planes,
|
||||
out_planes,
|
||||
kernel_size,
|
||||
stride,
|
||||
padding,
|
||||
groups=groups,
|
||||
bias=False,
|
||||
),
|
||||
nn.BatchNorm2d(out_planes),
|
||||
nn.ReLU6(inplace=True),
|
||||
)
|
||||
|
||||
|
||||
class InvertedResidual(nn.Module):
|
||||
|
||||
def __init__(self, inp, oup, stride, expand_ratio):
|
||||
super().__init__()
|
||||
self.stride = stride
|
||||
assert stride in [1, 2]
|
||||
|
||||
hidden_dim = int(round(inp * expand_ratio))
|
||||
self.use_res_connect = self.stride == 1 and inp == oup
|
||||
|
||||
layers = []
|
||||
if expand_ratio != 1:
|
||||
# pw
|
||||
layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1))
|
||||
layers.extend(
|
||||
[
|
||||
# dw
|
||||
ConvBNReLU(
|
||||
hidden_dim, hidden_dim, stride=stride, groups=hidden_dim
|
||||
),
|
||||
# pw-linear
|
||||
nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
|
||||
nn.BatchNorm2d(oup),
|
||||
]
|
||||
)
|
||||
self.conv = nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
if self.use_res_connect:
|
||||
return x + self.conv(x)
|
||||
else:
|
||||
return self.conv(x)
|
||||
|
||||
|
||||
class MobileNetV2(Backbone):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
width_mult=1.0,
|
||||
inverted_residual_setting=None,
|
||||
round_nearest=8,
|
||||
block=None,
|
||||
):
|
||||
"""
|
||||
MobileNet V2.
|
||||
|
||||
Args:
|
||||
width_mult (float): Width multiplier - adjusts number of channels in each layer by this amount
|
||||
inverted_residual_setting: Network structure
|
||||
round_nearest (int): Round the number of channels in each layer to be a multiple of this number
|
||||
Set to 1 to turn off rounding
|
||||
block: Module specifying inverted residual building block for mobilenet
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
if block is None:
|
||||
block = InvertedResidual
|
||||
input_channel = 32
|
||||
last_channel = 1280
|
||||
|
||||
if inverted_residual_setting is None:
|
||||
inverted_residual_setting = [
|
||||
# t, c, n, s
|
||||
[1, 16, 1, 1],
|
||||
[6, 24, 2, 2],
|
||||
[6, 32, 3, 2],
|
||||
[6, 64, 4, 2],
|
||||
[6, 96, 3, 1],
|
||||
[6, 160, 3, 2],
|
||||
[6, 320, 1, 1],
|
||||
]
|
||||
|
||||
# only check the first element, assuming user knows t,c,n,s are required
|
||||
if (
|
||||
len(inverted_residual_setting) == 0
|
||||
or len(inverted_residual_setting[0]) != 4
|
||||
):
|
||||
raise ValueError(
|
||||
"inverted_residual_setting should be non-empty "
|
||||
"or a 4-element list, got {}".
|
||||
format(inverted_residual_setting)
|
||||
)
|
||||
|
||||
# building first layer
|
||||
input_channel = _make_divisible(
|
||||
input_channel * width_mult, round_nearest
|
||||
)
|
||||
self.last_channel = _make_divisible(
|
||||
last_channel * max(1.0, width_mult), round_nearest
|
||||
)
|
||||
features = [ConvBNReLU(3, input_channel, stride=2)]
|
||||
# building inverted residual blocks
|
||||
for t, c, n, s in inverted_residual_setting:
|
||||
output_channel = _make_divisible(c * width_mult, round_nearest)
|
||||
for i in range(n):
|
||||
stride = s if i == 0 else 1
|
||||
features.append(
|
||||
block(
|
||||
input_channel, output_channel, stride, expand_ratio=t
|
||||
)
|
||||
)
|
||||
input_channel = output_channel
|
||||
# building last several layers
|
||||
features.append(
|
||||
ConvBNReLU(input_channel, self.last_channel, kernel_size=1)
|
||||
)
|
||||
# make it nn.Sequential
|
||||
self.features = nn.Sequential(*features)
|
||||
|
||||
self._out_features = self.last_channel
|
||||
|
||||
# weight initialization
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
nn.init.kaiming_normal_(m.weight, mode="fan_out")
|
||||
if m.bias is not None:
|
||||
nn.init.zeros_(m.bias)
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
nn.init.ones_(m.weight)
|
||||
nn.init.zeros_(m.bias)
|
||||
elif isinstance(m, nn.Linear):
|
||||
nn.init.normal_(m.weight, 0, 0.01)
|
||||
nn.init.zeros_(m.bias)
|
||||
|
||||
def _forward_impl(self, x):
|
||||
# This exists since TorchScript doesn't support inheritance, so the superclass method
|
||||
# (this one) needs to have a name other than `forward` that can be accessed in a subclass
|
||||
x = self.features(x)
|
||||
x = x.mean([2, 3])
|
||||
return x
|
||||
|
||||
def forward(self, x):
|
||||
return self._forward_impl(x)
|
||||
|
||||
|
||||
def init_pretrained_weights(model, model_url):
|
||||
"""Initializes model with pretrained weights.
|
||||
|
||||
Layers that don't match with pretrained layers in name or size are kept unchanged.
|
||||
"""
|
||||
if model_url is None:
|
||||
import warnings
|
||||
|
||||
warnings.warn(
|
||||
"ImageNet pretrained weights are unavailable for this model"
|
||||
)
|
||||
return
|
||||
pretrain_dict = model_zoo.load_url(model_url)
|
||||
model_dict = model.state_dict()
|
||||
pretrain_dict = {
|
||||
k: v
|
||||
for k, v in pretrain_dict.items()
|
||||
if k in model_dict and model_dict[k].size() == v.size()
|
||||
}
|
||||
model_dict.update(pretrain_dict)
|
||||
model.load_state_dict(model_dict)
|
||||
|
||||
|
||||
@BACKBONE_REGISTRY.register()
|
||||
def mobilenetv2(pretrained=True, **kwargs):
|
||||
model = MobileNetV2(**kwargs)
|
||||
if pretrained:
|
||||
init_pretrained_weights(model, model_urls["mobilenet_v2"])
|
||||
return model
|
||||
Reference in New Issue
Block a user