@@ -118,16 +118,28 @@ def __init__(self, stage, output_branches, c):
118118 nn .Upsample (scale_factor = (2.0 ** (branch_number - branch_output_number )), mode = 'nearest' ),
119119 ))
120120 elif branch_number < branch_output_number :
121- downsampling_fusion = []
122- for _ in range (branch_output_number - branch_number - 1 ):
123- downsampling_fusion .append (nn .Sequential (
124- nn .Conv2d (c * (2 ** branch_number ), c * (2 ** branch_number ), kernel_size = 3 , stride = 2 ,
125- padding = 1 ,
126- bias = False ),
127- nn .BatchNorm2d (c * (2 ** branch_number ), eps = 1e-05 , momentum = 0.1 , affine = True ,
128- track_running_stats = True ),
121+ downsampling_fusion = [
122+ nn .Sequential (
123+ nn .Conv2d (
124+ c * (2 ** branch_number ),
125+ c * (2 ** branch_number ),
126+ kernel_size = 3 ,
127+ stride = 2 ,
128+ padding = 1 ,
129+ bias = False ,
130+ ),
131+ nn .BatchNorm2d (
132+ c * (2 ** branch_number ),
133+ eps = 1e-05 ,
134+ momentum = 0.1 ,
135+ affine = True ,
136+ track_running_stats = True ,
137+ ),
129138 nn .ReLU (inplace = True ),
130- ))
139+ )
140+ for _ in range (branch_output_number - branch_number - 1 )
141+ ]
142+
131143 downsampling_fusion .append (nn .Sequential (
132144 nn .Conv2d (c * (2 ** branch_number ), c * (2 ** branch_output_number ), kernel_size = 3 ,
133145 stride = 2 , padding = 1 ,
@@ -224,15 +236,17 @@ def __init__(self, c=48, num_blocks=[1, 4, 3], num_classes=1000):
224236
225237 # Classifier (extra module if want to use for classification):
226238 # pool, reduce dimensionality, flatten, connect to linear layer for classification:
227- out_channels = sum ([ c * 2 ** i for i in range (len (num_blocks )+ 1 )]) # total output channels of HRNetV2
239+ out_channels = sum (c * 2 ** i for i in range (len (num_blocks )+ 1 ))
228240 pool_feature_map = 8
229241 self .bn_classifier = nn .Sequential (
230242 nn .Conv2d (out_channels , out_channels // 4 , kernel_size = 1 , bias = False ),
231- nn .BatchNorm2d (out_channels // 4 , eps = 1e-05 , affine = True , track_running_stats = True ),
243+ nn .BatchNorm2d (
244+ out_channels // 4 , eps = 1e-05 , affine = True , track_running_stats = True
245+ ),
232246 nn .ReLU (inplace = True ),
233247 nn .AdaptiveAvgPool2d (pool_feature_map ),
234248 nn .Flatten (),
235- nn .Linear (pool_feature_map * pool_feature_map * (out_channels // 4 ), num_classes ),
249+ nn .Linear (pool_feature_map ** 2 * (out_channels // 4 ), num_classes ),
236250 )
237251
238252 @staticmethod
0 commit comments