Spaces:
Running
Running
| # BSD 3-Clause License | |
| # Copyright (c) Soumith Chintala 2016, | |
| # All rights reserved. | |
| # Redistribution and use in source and binary forms, with or without | |
| # modification, are permitted provided that the following conditions are met: | |
| # * Redistributions of source code must retain the above copyright notice, this | |
| # list of conditions and the following disclaimer. | |
| # * Redistributions in binary form must reproduce the above copyright notice, | |
| # this list of conditions and the following disclaimer in the documentation | |
| # and/or other materials provided with the distribution. | |
| # * Neither the name of the copyright holder nor the names of its | |
| # contributors may be used to endorse or promote products derived from | |
| # this software without specific prior written permission. | |
| # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" | |
| # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE | |
| # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE | |
| # DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE | |
| # FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL | |
| # DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR | |
| # SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER | |
| # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, | |
| # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE | |
| # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. | |
| # Adaptation of the PyTorch torchvision MobileNetV2 without a classifier. | |
| # See source here: https://pytorch.org/vision/0.8/_modules/torchvision/models/mobilenet.html#mobilenet_v2 | |
| from torch import nn | |
| 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, norm_layer=None | |
| ): | |
| padding = (kernel_size - 1) // 2 | |
| if norm_layer is None: | |
| norm_layer = nn.BatchNorm2d | |
| super(ConvBNReLU, self).__init__( | |
| nn.Conv2d( | |
| in_planes, | |
| out_planes, | |
| kernel_size, | |
| stride, | |
| padding, | |
| groups=groups, | |
| bias=False, | |
| ), | |
| norm_layer(out_planes), | |
| nn.ReLU6(inplace=True), | |
| ) | |
| class InvertedResidual(nn.Module): | |
| def __init__(self, inp, oup, stride, expand_ratio, norm_layer=None): | |
| super(InvertedResidual, self).__init__() | |
| self.stride = stride | |
| assert stride in [1, 2] | |
| if norm_layer is None: | |
| norm_layer = nn.BatchNorm2d | |
| 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, norm_layer=norm_layer) | |
| ) | |
| layers.extend( | |
| [ | |
| # dw | |
| ConvBNReLU( | |
| hidden_dim, | |
| hidden_dim, | |
| stride=stride, | |
| groups=hidden_dim, | |
| norm_layer=norm_layer, | |
| ), | |
| # pw-linear | |
| nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False), | |
| norm_layer(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(nn.Module): | |
| def __init__( | |
| self, | |
| embed_dim=1028, | |
| width_mult=1.0, | |
| inverted_residual_setting=None, | |
| round_nearest=8, | |
| block=None, | |
| norm_layer=None, | |
| ): | |
| """ | |
| MobileNet V2 main class | |
| Args: | |
| embed_dim (int): Number of channels in the final output. | |
| 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 | |
| norm_layer: Module specifying the normalization layer to use | |
| """ | |
| super(MobileNetV2, self).__init__() | |
| if block is None: | |
| block = InvertedResidual | |
| if norm_layer is None: | |
| norm_layer = nn.BatchNorm2d | |
| input_channel = 32 | |
| last_channel = embed_dim / width_mult | |
| 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, norm_layer=norm_layer)] | |
| # 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, | |
| norm_layer=norm_layer, | |
| ) | |
| ) | |
| input_channel = output_channel | |
| # building last several layers | |
| features.append( | |
| ConvBNReLU( | |
| input_channel, self.last_channel, kernel_size=1, norm_layer=norm_layer | |
| ) | |
| ) | |
| # make it nn.Sequential | |
| self.features = nn.Sequential(*features) | |
| # 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.GroupNorm)): | |
| 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 | |
| return self.features(x) | |
| # return the features directly, no classifier or pooling | |
| def forward(self, x): | |
| return self._forward_impl(x) | |