@@ -275,7 +275,7 @@ class PoseHighResolutionNet(nn.Module):
275
275
276
276
def __init__ (self , cfg , ** kwargs ):
277
277
self .inplanes = 64
278
- extra = cfg . MODEL . EXTRA
278
+ extra = cfg [ ' MODEL' ][ ' EXTRA' ]
279
279
super (PoseHighResolutionNet , self ).__init__ ()
280
280
281
281
# stem net
@@ -288,7 +288,7 @@ def __init__(self, cfg, **kwargs):
288
288
self .relu = nn .ReLU (inplace = True )
289
289
self .layer1 = self ._make_layer (Bottleneck , 64 , 4 )
290
290
291
- self .stage2_cfg = cfg [ 'MODEL' ][ 'EXTRA' ] ['STAGE2' ]
291
+ self .stage2_cfg = extra ['STAGE2' ]
292
292
num_channels = self .stage2_cfg ['NUM_CHANNELS' ]
293
293
block = blocks_dict [self .stage2_cfg ['BLOCK' ]]
294
294
num_channels = [
@@ -298,7 +298,7 @@ def __init__(self, cfg, **kwargs):
298
298
self .stage2 , pre_stage_channels = self ._make_stage (
299
299
self .stage2_cfg , num_channels )
300
300
301
- self .stage3_cfg = cfg [ 'MODEL' ][ 'EXTRA' ] ['STAGE3' ]
301
+ self .stage3_cfg = extra ['STAGE3' ]
302
302
num_channels = self .stage3_cfg ['NUM_CHANNELS' ]
303
303
block = blocks_dict [self .stage3_cfg ['BLOCK' ]]
304
304
num_channels = [
@@ -309,7 +309,7 @@ def __init__(self, cfg, **kwargs):
309
309
self .stage3 , pre_stage_channels = self ._make_stage (
310
310
self .stage3_cfg , num_channels )
311
311
312
- self .stage4_cfg = cfg [ 'MODEL' ][ 'EXTRA' ] ['STAGE4' ]
312
+ self .stage4_cfg = extra ['STAGE4' ]
313
313
num_channels = self .stage4_cfg ['NUM_CHANNELS' ]
314
314
block = blocks_dict [self .stage4_cfg ['BLOCK' ]]
315
315
num_channels = [
@@ -322,13 +322,13 @@ def __init__(self, cfg, **kwargs):
322
322
323
323
self .final_layer = nn .Conv2d (
324
324
in_channels = pre_stage_channels [0 ],
325
- out_channels = cfg . MODEL . NUM_JOINTS ,
326
- kernel_size = extra . FINAL_CONV_KERNEL ,
325
+ out_channels = cfg [ ' MODEL' ][ ' NUM_JOINTS' ] ,
326
+ kernel_size = extra [ ' FINAL_CONV_KERNEL' ] ,
327
327
stride = 1 ,
328
- padding = 1 if extra . FINAL_CONV_KERNEL == 3 else 0
328
+ padding = 1 if extra [ ' FINAL_CONV_KERNEL' ] == 3 else 0
329
329
)
330
330
331
- self .pretrained_layers = cfg [ 'MODEL' ][ 'EXTRA' ] ['PRETRAINED_LAYERS' ]
331
+ self .pretrained_layers = extra ['PRETRAINED_LAYERS' ]
332
332
333
333
def _make_transition_layer (
334
334
self , num_channels_pre_layer , num_channels_cur_layer ):
@@ -495,7 +495,7 @@ def init_weights(self, pretrained=''):
495
495
def get_pose_net (cfg , is_train , ** kwargs ):
496
496
model = PoseHighResolutionNet (cfg , ** kwargs )
497
497
498
- if is_train and cfg . MODEL . INIT_WEIGHTS :
499
- model .init_weights (cfg . MODEL . PRETRAINED )
498
+ if is_train and cfg [ ' MODEL' ][ ' INIT_WEIGHTS' ] :
499
+ model .init_weights (cfg [ ' MODEL' ][ ' PRETRAINED' ] )
500
500
501
501
return model
0 commit comments