fromfunctoolsimportpartialfromtypingimportAny,Callable,List,Optionalimporttorchfromtorchimportnn,Tensorfrom..ops.miscimportConv2dNormActivationfrom..transforms._presetsimportImageClassificationfrom..utilsimport_log_api_usage_oncefrom._apiimportregister_model,Weights,WeightsEnumfrom._metaimport_IMAGENET_CATEGORIESfrom._utilsimport_make_divisible,_ovewrite_named_param,handle_legacy_interface__all__=["MobileNetV2","MobileNet_V2_Weights","mobilenet_v2"]# necessary for backwards compatibilityclassInvertedResidual(nn.Module):def__init__(self,inp:int,oup:int,stride:int,expand_ratio:int,norm_layer:Optional[Callable[...,nn.Module]]=None)->None:super().__init__()self.stride=strideifstridenotin[1,2]:raiseValueError(f"stride should be 1 or 2 instead of {stride}")ifnorm_layerisNone:norm_layer=nn.BatchNorm2dhidden_dim=int(round(inp*expand_ratio))self.use_res_connect=self.stride==1andinp==ouplayers:List[nn.Module]=[]ifexpand_ratio!=1:# pwlayers.append(Conv2dNormActivation(inp,hidden_dim,kernel_size=1,norm_layer=norm_layer,activation_layer=nn.ReLU6))layers.extend([# dwConv2dNormActivation(hidden_dim,hidden_dim,stride=stride,groups=hidden_dim,norm_layer=norm_layer,activation_layer=nn.ReLU6,),# pw-linearnn.Conv2d(hidden_dim,oup,1,1,0,bias=False),norm_layer(oup),])self.conv=nn.Sequential(*layers)self.out_channels=oupself._is_cn=stride>1defforward(self,x:Tensor)->Tensor:ifself.use_res_connect:returnx+self.conv(x)else:returnself.conv(x)classMobileNetV2(nn.Module):def__init__(self,num_classes:int=1000,width_mult:float=1.0,inverted_residual_setting:Optional[List[List[int]]]=None,round_nearest:int=8,block:Optional[Callable[...,nn.Module]]=None,norm_layer:Optional[Callable[...,nn.Module]]=None,dropout:float=0.2,)->None:""" MobileNet V2 main class Args: num_classes (int): Number of classes 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 dropout (float): The droupout probability """super().__init__()_log_api_usage_once(self)ifblockisNone:block=InvertedResidualifnorm_layerisNone:norm_layer=nn.BatchNorm2dinput_channel=32last_channel=1280ifinverted_residual_settingisNone: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 requirediflen(inverted_residual_setting)==0orlen(inverted_residual_setting[0])!=4:raiseValueError(f"inverted_residual_setting should be non-empty or a 4-element list, got {inverted_residual_setting}")# building first layerinput_channel=_make_divisible(input_channel*width_mult,round_nearest)self.last_channel=_make_divisible(last_channel*max(1.0,width_mult),round_nearest)features:List[nn.Module]=[Conv2dNormActivation(3,input_channel,stride=2,norm_layer=norm_layer,activation_layer=nn.ReLU6)]# building inverted residual blocksfort,c,n,sininverted_residual_setting:output_channel=_make_divisible(c*width_mult,round_nearest)foriinrange(n):stride=sifi==0else1features.append(block(input_channel,output_channel,stride,expand_ratio=t,norm_layer=norm_layer))input_channel=output_channel# building last several layersfeatures.append(Conv2dNormActivation(input_channel,self.last_channel,kernel_size=1,norm_layer=norm_layer,activation_layer=nn.ReLU6))# make it nn.Sequentialself.features=nn.Sequential(*features)# building classifierself.classifier=nn.Sequential(nn.Dropout(p=dropout),nn.Linear(self.last_channel,num_classes),)# weight initializationforminself.modules():ifisinstance(m,nn.Conv2d):nn.init.kaiming_normal_(m.weight,mode="fan_out")ifm.biasisnotNone:nn.init.zeros_(m.bias)elifisinstance(m,(nn.BatchNorm2d,nn.GroupNorm)):nn.init.ones_(m.weight)nn.init.zeros_(m.bias)elifisinstance(m,nn.Linear):nn.init.normal_(m.weight,0,0.01)nn.init.zeros_(m.bias)def_forward_impl(self,x:Tensor)->Tensor:# 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 subclassx=self.features(x)# Cannot use "squeeze" as batch-size can be 1x=nn.functional.adaptive_avg_pool2d(x,(1,1))x=torch.flatten(x,1)x=self.classifier(x)returnxdefforward(self,x:Tensor)->Tensor:returnself._forward_impl(x)_COMMON_META={"num_params":3504872,"min_size":(1,1),"categories":_IMAGENET_CATEGORIES,}
[docs]classMobileNet_V2_Weights(WeightsEnum):IMAGENET1K_V1=Weights(url="https://download.pytorch.org/models/mobilenet_v2-b0353104.pth",transforms=partial(ImageClassification,crop_size=224),meta={**_COMMON_META,"recipe":"https://github.com/pytorch/vision/tree/main/references/classification#mobilenetv2","_metrics":{"ImageNet-1K":{"acc@1":71.878,"acc@5":90.286,}},"_ops":0.301,"_file_size":13.555,"_docs":"""These weights reproduce closely the results of the paper using a simple training recipe.""",},)IMAGENET1K_V2=Weights(url="https://download.pytorch.org/models/mobilenet_v2-7ebf99e0.pth",transforms=partial(ImageClassification,crop_size=224,resize_size=232),meta={**_COMMON_META,"recipe":"https://github.com/pytorch/vision/issues/3995#new-recipe-with-reg-tuning","_metrics":{"ImageNet-1K":{"acc@1":72.154,"acc@5":90.822,}},"_ops":0.301,"_file_size":13.598,"_docs":""" These weights improve upon the results of the original paper by using a modified version of TorchVision's `new training recipe <https://pytorch.org/blog/how-to-train-state-of-the-art-models-using-torchvision-latest-primitives/>`_. """,},)DEFAULT=IMAGENET1K_V2
[docs]@register_model()@handle_legacy_interface(weights=("pretrained",MobileNet_V2_Weights.IMAGENET1K_V1))defmobilenet_v2(*,weights:Optional[MobileNet_V2_Weights]=None,progress:bool=True,**kwargs:Any)->MobileNetV2:"""MobileNetV2 architecture from the `MobileNetV2: Inverted Residuals and Linear Bottlenecks <https://arxiv.org/abs/1801.04381>`_ paper. Args: weights (:class:`~torchvision.models.MobileNet_V2_Weights`, optional): The pretrained weights to use. See :class:`~torchvision.models.MobileNet_V2_Weights` below for more details, and possible values. By default, no pre-trained weights are used. progress (bool, optional): If True, displays a progress bar of the download to stderr. Default is True. **kwargs: parameters passed to the ``torchvision.models.mobilenetv2.MobileNetV2`` base class. Please refer to the `source code <https://github.com/pytorch/vision/blob/main/torchvision/models/mobilenetv2.py>`_ for more details about this class. .. autoclass:: torchvision.models.MobileNet_V2_Weights :members: """weights=MobileNet_V2_Weights.verify(weights)ifweightsisnotNone:_ovewrite_named_param(kwargs,"num_classes",len(weights.meta["categories"]))model=MobileNetV2(**kwargs)ifweightsisnotNone:model.load_state_dict(weights.get_state_dict(progress=progress,check_hash=True))returnmodel
Docs
Access comprehensive developer documentation for PyTorch
To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. As the current maintainers of this site, Facebook’s Cookies Policy applies. Learn more, including about available controls: Cookies Policy.