a dA @s~ddlmZddlmZmZmZmZddlZddlm Z ddlm Z ddl m Z ddl mZdd lmZmZmZdd lmZdd lmZmZdd l mZmZmZmZd dlmZmZmZgdZ Gddde j!Z"Gddde j#Z$ee%ee%eee&e&ee$dddZ'deddddZ(GdddeZ)GdddeZ*Gd d!d!eZ+Gd"d#d#eZ,ed$d%ed&d'd(fd)dd*d+d,eee)efe&e&ee$d-d.d/Z-ed0d%ed&d1d(fd)dd*d+d,eee*efe&e&ee$d-d2d3Z.ed4d%ed&d5d(fd)dd*d+d,eee+efe&e&ee$d-d6d7Z/ed8d%ed&d9d(fd)dd*d+d,eee,efe&e&ee$d-d:d;Z0dS)<)partial)AnyListOptionalUnionN)Tensor) shufflenetv2)ImageClassification)register_modelWeights WeightsEnum)_IMAGENET_CATEGORIES)_ovewrite_named_paramhandle_legacy_interface)ShuffleNet_V2_X0_5_WeightsShuffleNet_V2_X1_0_WeightsShuffleNet_V2_X1_5_WeightsShuffleNet_V2_X2_0_Weights) _fuse_modules _replace_reluquantize_model) QuantizableShuffleNetV2#ShuffleNet_V2_X0_5_QuantizedWeights#ShuffleNet_V2_X1_0_QuantizedWeights#ShuffleNet_V2_X1_5_QuantizedWeights#ShuffleNet_V2_X2_0_QuantizedWeightsshufflenet_v2_x0_5shufflenet_v2_x1_0shufflenet_v2_x1_5shufflenet_v2_x2_0cs6eZdZeeddfdd ZeedddZZS)QuantizableInvertedResidualNargskwargsreturncs"tj|i|tj|_dSN)super__init__nnZ quantizedZFloatFunctionalcatselfr%r& __class__u/var/www/html/stable-diffusion-webui/venv/lib/python3.9/site-packages/torchvision/models/quantization/shufflenetv2.pyr*$sz$QuantizableInvertedResidual.__init__xr'cCsh|jdkr8|jddd\}}|jj|||gdd}n |jj||||gdd}t|d}|S)Nrr )Zdim)Zstridechunkr,branch2branch1rZchannel_shuffle)r.r4x1Zx2outr1r1r2forward(s    z#QuantizableInvertedResidual.forward)__name__ __module__ __qualname__rr*rr: __classcell__r1r1r/r2r##sr#csLeZdZeeddfdd ZeedddZd eeddd d Z Z S) rNr$cs6tj|dti|tjj|_tjj|_ dS)NZinverted_residual) r)r*r#torchZaoZ quantizationZ QuantStubquantZ DeQuantStubdequantr-r/r1r2r*6sz QuantizableShuffleNetV2.__init__r3cCs"||}||}||}|Sr()r@Z _forward_implrA)r.r4r1r1r2r:;s   zQuantizableShuffleNetV2.forward)is_qatr'cCs|jD].\}}|dvr |dur t|gdg|ddq |D]f}t|turBt|jjdkrt|jddggd g|ddt|jgdd d ggd g|ddqBdS) aBFuse conv/bn/relu modules in shufflenetv2 model Fuse conv+bn+relu/ conv+relu/conv+bn modules to prepare for quantization. Model is modified in place. .. note:: Note that this operation does not change numerics and the model after modification is in floating point )Zconv1Zconv5N)012T)ZinplacerrCrD)rE34rFrG)567) Z_modulesitemsrmodulestyper#lenr7r6)r.rBnamemr1r1r2 fuse_modelAs   z"QuantizableShuffleNetV2.fuse_model)N) r;r<r=rr*rr:rboolrQr>r1r1r/r2r4sr)stages_repeatsstages_out_channelsweightsprogressquantizer&r'cKs|dur:t|dt|jdd|jvr:t|d|jd|dd}t||fi|}t||rnt|||dur||j|d|S)NZ num_classes categoriesbackendfbgemm)rV) rrNmetapoprrrZload_state_dictZget_state_dict)rSrTrUrVrWr&rYmodelr1r1r2 _shufflenetv2Zs    r^)rrrZzdhttps://github.com/pytorch/vision/tree/main/references/classification#post-training-quantized-modelsz These weights were produced by doing Post Training Quantization (eager mode) on top of the unquantized weights listed below. )Zmin_sizerXrYrecipeZ_docsc @sFeZdZedeeddiedejddddid d d d Z e Z d S)rzShttps://download.pytorch.org/models/quantized/shufflenetv2_x0.5_fbgemm-00845098.pth crop_sizei ImageNet-1Kg#~jL@gRS@zacc@1zacc@5g{Gz?gjt? num_params unquantized_metrics_ops _file_sizeurlZ transformsr[N) r;r<r=r rr _COMMON_METAr IMAGENET1K_V1IMAGENET1K_FBGEMM_V1DEFAULTr1r1r1r2rs$ rc @sFeZdZedeeddiedejddddid d d d Z e Z d S)rzQhttps://download.pytorch.org/models/quantized/shufflenetv2_x1_fbgemm-db332c57.pthr`rai"rcgףp= Q@gh|?U@rdg(\?gy&1@rerkN) r;r<r=r rr rmrrnrorpr1r1r1r2rs$ rc @sJeZdZedeedddieddejddd d id d d dZ e Z dS)rzShttps://download.pytorch.org/models/quantized/shufflenetv2_x1_5_fbgemm-d7401f05.pthr`rbZ resize_size+https://github.com/pytorch/vision/pull/5906iv5rcgSR@g̬V@rdgl?gK7A` @r_rfrgrhrirjrkN) r;r<r=r rr rmrrnrorpr1r1r1r2rs& rc @sJeZdZedeedddieddejddd d id d d dZ e Z dS)rzShttps://download.pytorch.org/models/quantized/shufflenetv2_x2_0_fbgemm-5cac526c.pthr`rqrrrsiprcg-R@gZd;W@rdg-?g|?5@rtrkN) r;r<r=r rr rmrrnrorpr1r1r1r2rs& rZquantized_shufflenet_v2_x0_5)rOZ pretrainedcCs|ddrtjStjSNrWF)getrrorrnr&r1r1r2s rx)rUTFrUrVrW)rUrVrWr&r'cKs4|rtnt|}tgdgdf|||d|S)aQ Constructs a ShuffleNetV2 with 0.5x output channels, as described in `ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design `__. .. note:: Note that ``quantize = True`` returns a quantized model with 8 bit weights. Quantized models only support inference and run on CPUs. GPU inference is not yet supported. Args: weights (:class:`~torchvision.models.quantization.ShuffleNet_V2_X0_5_QuantizedWeights` or :class:`~torchvision.models.ShuffleNet_V2_X0_5_Weights`, optional): The pretrained weights for the model. See :class:`~torchvision.models.quantization.ShuffleNet_V2_X0_5_QuantizedWeights` 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. quantize (bool, optional): If True, return a quantized version of the model. Default is False. **kwargs: parameters passed to the ``torchvision.models.quantization.ShuffleNet_V2_X0_5_QuantizedWeights`` base class. Please refer to the `source code `_ for more details about this class. .. autoclass:: torchvision.models.quantization.ShuffleNet_V2_X0_5_QuantizedWeights :members: .. autoclass:: torchvision.models.ShuffleNet_V2_X0_5_Weights :members: :noindex: r{)0`ry)rrverifyr^rUrVrWr&r1r1r2rs0 rZquantized_shufflenet_v2_x1_0cCs|ddrtjStjSru)rvrrorrnrwr1r1r2rxs cKs4|rtnt|}tgdgdf|||d|S)aQ Constructs a ShuffleNetV2 with 1.0x output channels, as described in `ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design `__. .. note:: Note that ``quantize = True`` returns a quantized model with 8 bit weights. Quantized models only support inference and run on CPUs. GPU inference is not yet supported. Args: weights (:class:`~torchvision.models.quantization.ShuffleNet_V2_X1_0_QuantizedWeights` or :class:`~torchvision.models.ShuffleNet_V2_X1_0_Weights`, optional): The pretrained weights for the model. See :class:`~torchvision.models.quantization.ShuffleNet_V2_X1_0_QuantizedWeights` 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. quantize (bool, optional): If True, return a quantized version of the model. Default is False. **kwargs: parameters passed to the ``torchvision.models.quantization.ShuffleNet_V2_X1_0_QuantizedWeights`` base class. Please refer to the `source code `_ for more details about this class. .. autoclass:: torchvision.models.quantization.ShuffleNet_V2_X1_0_QuantizedWeights :members: .. autoclass:: torchvision.models.ShuffleNet_V2_X1_0_Weights :members: :noindex: rz)r}trqirry)rrrr^rr1r1r2r s0 r Zquantized_shufflenet_v2_x1_5cCs|ddrtjStjSru)rvrrorrnrwr1r1r2rxFs cKs4|rtnt|}tgdgdf|||d|S)aQ Constructs a ShuffleNetV2 with 1.5x output channels, as described in `ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design `__. .. note:: Note that ``quantize = True`` returns a quantized model with 8 bit weights. Quantized models only support inference and run on CPUs. GPU inference is not yet supported. Args: weights (:class:`~torchvision.models.quantization.ShuffleNet_V2_X1_5_QuantizedWeights` or :class:`~torchvision.models.ShuffleNet_V2_X1_5_Weights`, optional): The pretrained weights for the model. See :class:`~torchvision.models.quantization.ShuffleNet_V2_X1_5_QuantizedWeights` 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. quantize (bool, optional): If True, return a quantized version of the model. Default is False. **kwargs: parameters passed to the ``torchvision.models.quantization.ShuffleNet_V2_X1_5_QuantizedWeights`` base class. Please refer to the `source code `_ for more details about this class. .. autoclass:: torchvision.models.quantization.ShuffleNet_V2_X1_5_QuantizedWeights :members: .. autoclass:: torchvision.models.ShuffleNet_V2_X1_5_Weights :members: :noindex: rz)r}i`irry)rrrr^rr1r1r2r!Bs0 r!Zquantized_shufflenet_v2_x2_0cCs|ddrtjStjSru)rvrrorrnrwr1r1r2rx|s cKs4|rtnt|}tgdgdf|||d|S)aQ Constructs a ShuffleNetV2 with 2.0x output channels, as described in `ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design `__. .. note:: Note that ``quantize = True`` returns a quantized model with 8 bit weights. Quantized models only support inference and run on CPUs. GPU inference is not yet supported. Args: weights (:class:`~torchvision.models.quantization.ShuffleNet_V2_X2_0_QuantizedWeights` or :class:`~torchvision.models.ShuffleNet_V2_X2_0_Weights`, optional): The pretrained weights for the model. See :class:`~torchvision.models.quantization.ShuffleNet_V2_X2_0_QuantizedWeights` 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. quantize (bool, optional): If True, return a quantized version of the model. Default is False. **kwargs: parameters passed to the ``torchvision.models.quantization.ShuffleNet_V2_X2_0_QuantizedWeights`` base class. Please refer to the `source code `_ for more details about this class. .. autoclass:: torchvision.models.quantization.ShuffleNet_V2_X2_0_QuantizedWeights :members: .. autoclass:: torchvision.models.ShuffleNet_V2_X2_0_Weights :members: :noindex: rz)r}iiiry)rrrr^rr1r1r2r"xs0 r")1 functoolsrtypingrrrrr?Ztorch.nnr+rZtorchvision.modelsrZtransforms._presetsr Z_apir r rZ_metar_utilsrrrrrrutilsrrr__all__ZInvertedResidualr#Z ShuffleNetV2rintrRr^rmrrrrrr r!r"r1r1r1r2s       '   - - -