B
|
ɧ_éD ã @ sÀ d Z ddlZddlZddlZddlmZ ddlmZmZmZm Z ddl
|
m Z m Z ddl mZ dZd d
|
d D Zd d Zdd
|
e ¡ D ZG dd dejZG dd dejZe ¡ dd ZdS )zResNeSt modelsé N)Únn)ÚIBNÚ Non_localÚ SplAtConv2dÚget_norm)Ú!get_unexpected_parameters_messageÚ get_missing_parameters_messageé )ÚBACKBONE_REGISTRYz9https://hangzh.s3.amazonaws.com/encoding/models/{}-{}.pthc C s i | ]\}}||qS © r )Ú.0ZchecksumÚnamer r ú ./modeling/backbones/resnest.pyú
|
<dictcomp> s r ))Z528c19caZ resnest50)Z22405ba7Z
|
resnest101)Z75117900Z
|
resnest200)Z0cc87c48Z
|
resnest269c C s( | t krtdj| dt | d d
S )Nz-Pretrained model for {name} is not available.)r é )Ú _model_sha256Ú
|
ValueErrorÚformat)r r r r Ú
|
short_hash! s r c C s i | ]}t |t|¡|qS r )Ú _url_formatr r )r r r r r r ' s c s. e Zd ZdZdZd fdd Zd
|
d Z ZS ) Ú
|
BottleneckzResNet Bottleneck
|
é Fr Né@ ç c s t t| ¡ t||
|
d | }tj||ddd| _|rJt|||| _nt |||| _|| _
|
|| _ | or|dkpr|| _ | | _ | j rtjd|dd| _d}|dkrÈt||d|| | | d||||||d| _nh|rdd lm} |||d|| | | d|d
|
| _t |||| _n*tj||d|| | | dd | _t |||| _tj||d ddd| _t ||d || _|rxdd lm} || jj tjdd| _|| _ | | _ || _ d S )Ng P@r F)Ú kernel_sizeÚbiasé )Úpadding) r Ústrider ÚdilationÚgroupsr ÚradixÚrectifyÚ rectify_avgÚ
|
norm_layerÚ
|
num_splitsÚdropblock_probr )ÚRFConv2d)r r r r r r Ú average_mode)r r r r r r r )Úzeros_T)Úinplace) Úsuperr Ú__init__Úintr ÚConv2dÚconv1r Úbn1r r& r! ÚavdÚ avd_firstÚ AvgPool2dÚ avd_layerr Úconv2Úrfconvr' Úbn2Úconv3Úbn3Z torch.nn.initr) ÚweightÚReLUÚreluÚ
|
downsampler r )ÚselfÚinplanesÚplanesÚbn_normr% Úwith_ibnr r= r! Ú cardinalityÚbottleneck_widthr1 r2 r Úis_firstÚrectified_convr# r& Ú
|
last_gammaZ group_widthr' r) )Ú __class__r r r, 2 sZ
|
zBottleneck.__init__c C sð |}| |¡}| |¡}| jdkr,| |¡}| |¡}| jrL| jrL| |¡}| |¡}| j dkr|
|
|¡}| jdkr~| |¡}| |¡}| jr| js| |¡}| |¡}| |¡}| jdkrÆ| |¡}| jd k rÚ| |¡}||7 }| |¡}|S )Ng r )r/ r0 r& Z
|
dropblock1r< r1 r2 r4 r5 r! r7 Z
|
dropblock2r8 r9 Z
|
dropblock3r= )r> ÚxÚresidualÚoutr r r Úforwardk s2
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
zBottleneck.forward)Fr Nr r r FFr FFFr F)Ú__name__Ú
|
__module__Ú __qualname__Ú__doc__Ú expansionr, rL Ú __classcell__r r )rH r r , s 5r c s< e Zd ZdZd fdd Zdd d Zd d Zdd Z ZS )ÚResNesta ResNet Variants ResNest
|
Parameters
|
----------
|
block : Block
|
Class for the residual block. Options are BasicBlockV1, BottleneckV1.
|
layers : list of int
|
Numbers of layers in each block
|
classes : int, default 1000
|
Number of classification classes.
|
dilated : bool, default False
|
Applying dilation strategy to pretrained ResNet yielding a stride-8 model,
|
typically used in Semantic Segmentation.
|
norm_layer : object
|
Normalization layer used in backbone network (default: :class:`mxnet.gluon.nn.BatchNorm`;
|
for Synchronized Cross-GPU BachNormalization).
|
Reference:
|
- He, Kaiming, et al. "Deep residual learning for image recognition." Proceedings of the IEEE conference on computer vision and pattern recognition. 2016.
|
- Yu, Fisher, and Vladlen Koltun. "Multi-scale context aggregation by dilated convolutions."
|
r r Fç r c s, |
|
| _ | | _|r|d nd| _|| _|| _| | _|| _|| _t ¡ || _
|
|| _ |rhddl m } |}ntj}|rzd|ini }|r
|
t |d|fddddd |t|||tjd
|
d |||fddddd |t|||tjd
|
d |||d fddddd |¡| _n|dd dddd || _t|| j|| _tjd
|
d | _tjdddd | _| j|d|d d|||dd| _| j|d|d d|||d| _| s¢| dkrè| j|d|d d|||d|d | _| j|d|d d|||d|d | _ n| dkr8| j|d|d d|||d|d | _| j|d|d d|||d|d | _ n@| j|d|d d||||d| _| j|d|d |||||d| _ x~| ¡ D ]r}t |tjrÊ|j d |j d |j } |j!j" #dt$ %d| ¡¡ n(t |tj&r|j!j" 'd¡ |j(j" )¡ qW |r| *||||¡ ng | _+ | _, | _-| _.d S )Né r r )r' r( r r F)r r r r T)r* é )r r r )rB rE é )rB r é )rB r r& i )rB r& g @)r r )/rC rD r? Úavg_downrG r! r1 r2 r+ r, rF r# r6 r' r r. Ú
|
Sequentialr r; r/ r0 r< Ú MaxPool2dÚmaxpoolÚ _make_layerÚlayer1Úlayer2Úlayer3Úlayer4ÚmodulesÚ
|
isinstancer Ú out_channelsr: ÚdataÚnormal_ÚmathÚsqrtÚ BatchNorm2dÚfill_r Úzero_Ú_build_nonlocalÚNL_1_idxÚNL_2_idxÚNL_3_idxÚNL_4_idx) r> Ú last_striderA r% rB Úwith_nlÚblockÚlayersÚ
|
non_layersr! r rD Zdilatedr Ú deep_stemÚ
|
stem_widthrY rF r# r1 r2 Z
|
final_dropr& rG r' Z
|
conv_layerZ conv_kwargsÚmÚn)rH r r r, ¥ sp
|
|
|
|
|
&
|
|
|
|
zResNest.__init__ÚBNTc C sø d } |dks | j ||j krÈg } | jr|dkrH| tj||ddd¡ n| tjddddd¡ | tj| j ||j dddd¡ n"| tj| j ||j d|dd¡ | t|||j |¡ tj| } g } |dkrØd}|dksê|dkr0| || j |||||| | j | j
|
| j | j | j d|
|
| j| j| | jd¡ n^|d kr| || j |||||| | j | j
|
| j | j | j d|
|
| j| j| | jd¡ ntd
|
|¡||j | _ xRtd|D ]D}| || j ||||| j | j
|
| j | j | j || j| j| | jd ¡ q¦W tj| S ) Nr TF)r r Ú ceil_modeÚcount_include_pad)r r r i rU ) r= r! rC rD r1 r2 r rE rF r# r& rG r z => unknown dilation size: {})
|
r! rC rD r1 r2 r rF r# r& rG )r? rQ rY Úappendr r3 r. r rZ r! rC rD r1 r2 rF r# rG Ú RuntimeErrorr Úrange)r> rs r@ Úblocksr rA r% rB r r& rE r= Z down_layersrt Úir r r r] ð s`
|
|
zResNest._make_layerc s t fddt|d D ¡| _tfddt|d D | _t fddt|d D ¡| _tfddt|d D | _t fddt|d D ¡| _tfd
|
dt|d D | _ t fd dt|d D ¡| _
|
tfd dt|d D | _ d S )Nc s g | ]}t d qS )rX )r )r Ú_)rA r% r r ú
|
<listcomp>+ s z+ResNest._build_nonlocal.<locals>.<listcomp>r c s g | ]} d |d qS )r r r )r r )rt r r r , s c s g | ]}t d qS )i )r )r r )rA r% r r r . s r c s g | ]} d |d qS )r r )r r )rt r r r / s c s g | ]}t d qS )i )r )r r )rA r% r r r 1 s rU c s g | ]} d |d qS )rU r r )r r )rt r r r 2 s c s g | ]}t d qS )i )r )r r )rA r% r r r 4 s r c s g | ]} d |d qS )r r r )r r )rt r r r 5 s ) r Ú
|
ModuleListr ÚNL_1Úsortedrm ÚNL_2rn ÚNL_3ro ÚNL_4rp )r> rt ru rA r% r )rA rt r% r rl ) s zResNest._build_nonlocalc C s | |¡}| |¡}| |¡}| |¡}d}t| jdkrBdg| _xVtt| jD ]D}| j| |}|| j| krR|j\}}}}| j | |}|d7 }qRW d}t| j
|
dkr´dg| _
|
xVtt| j D ]D}| j | |}|| j
|
| krÄ|j\}}}}| j | |}|d7 }qÄW d} t| j dkr(dg| _ xZtt| jD ]H}| j| |}|| j | kr8|j\}}}}| j| |}| d7 } q8W d}
|
t| jdkr dg| _xZtt| jD ]H}| j| |}|| j|
|
kr°|j\}}}}| j|
|
|}|
|
d7 }
|
q°W |S )Nr éÿÿÿÿr )r/ r0 r< r\ Úlenrm r r^ Úshaper
rn r_ r ro r` r rp ra r ) r> rI Ú NL1_counterr r ÚCÚHÚWÚ NL2_counterÚ NL3_counterÚ NL4_counterr r r rL 7 sR
|
|
|
|
zResNest.forward)r r r Fr Fr FFFFFrT r F)r rz r Fr rT T) rM rN rO rP r, r] rl rL rR r r )rH r rS s D
|
8rS c C sJ | j jj}| j jj}| j jj}| j jj}| j jj}| j jj}| j jj}| j jj }ddddgddddgddddgdddd gd
|
| } d d dd gd d dd gd | }
|
ddddd
|
| } t
|
|||||t | |
|
d ddd| dddd} |rFt j jtdt| ddd} | j| dd}t t¡}|jr.| t|j¡ |jrF| t|j¡ | S )zi
|
Create a ResNest instance from config.
|
Returns:
|
ResNet: a :class:`ResNet` instance.
|
r r é é é é$ é é0 r )é2 ée éÈ i r rU )r r é r r TF)r! r rD rv rw rY r1 r2 Úresnest)ÚprogressÚ
|
check_hash)Ústrict)ÚMODELÚBACKBONEÚPRETRAINÚ LAST_STRIDEÚNORMZ
|
NORM_SPLITÚWITH_IBNÚWITH_SEÚWITH_NLÚDEPTHrS r ÚtorchÚhubÚload_state_dict_from_urlÚ
|
model_urlsÚstrÚload_state_dictÚloggingÚ getLoggerrM Ú missing_keysÚinfor Úunexpected_keysr )ÚcfgÚpretrainrq rA r% rB Úwith_serr ÚdepthÚnum_blocks_per_stageÚnl_layers_per_stagerw ÚmodelÚ
|
state_dictÚ incompatibleÚloggerr r r Úbuild_resnest_backboneh s6
|
|
|
|
|
|
|
|
2
|
|
rÀ )rP r± rg r« r rt r r r r Úutils.checkpointr r Úbuildr
|
r r r Úkeysr® ÚModuler rS ÚregisterrÀ r r r r Ú<module> s c Z
|
