update esnet.py

ppcls/arch/backbone/legendary_models/esnet.py

@@ -18,14 +18,13 @@ import paddle
 from paddle import ParamAttr, reshape, transpose, concat, split
 import paddle.nn as nn
 from paddle.nn import Conv2D, BatchNorm, Linear, Dropout
-from paddle.nn import AdaptiveAvgPool2D, MaxPool2D, AvgPool2D
+from paddle.nn import AdaptiveAvgPool2D, MaxPool2D
 from paddle.nn.initializer import KaimingNormal
 from paddle.regularizer import L2Decay
 
 from ppcls.arch.backbone.base.theseus_layer import TheseusLayer
 from ppcls.utils.save_load import load_dygraph_pretrain, load_dygraph_pretrain_from_url
 
-
 MODEL_URLS = {
     "ESNet_x0_25":
     "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/legendary_models/ESNet_x0_25_pretrained.pdparams",
@@ -60,8 +59,7 @@ def make_divisible(v, divisor=8, min_value=None):
 
 
 class ConvBNLayer(TheseusLayer):
-    def __init__(
-            self,
+    def __init__(self,
                  in_channels,
                  out_channels,
                  kernel_size,
@@ -125,9 +123,7 @@ class SEModule(TheseusLayer):
 
 
 class ESBlock1(TheseusLayer):
-    def __init__(self,
-                 in_channels,
-                 out_channels):
+    def __init__(self, in_channels, out_channels):
         super().__init__()
         self.pw_1_1 = ConvBNLayer(
             in_channels=in_channels // 2,
@@ -151,9 +147,7 @@ class ESBlock1(TheseusLayer):
 
     def forward(self, x):
         x1, x2 = split(
-            x,
-            num_or_sections=[x.shape[1] // 2, x.shape[1] // 2],
-            axis=1)
+            x, num_or_sections=[x.shape[1] // 2, x.shape[1] // 2], axis=1)
         x2 = self.pw_1_1(x2)
         x3 = self.dw_1(x2)
         x3 = concat([x2, x3], axis=1)
@@ -164,9 +158,7 @@ class ESBlock1(TheseusLayer):
 
 
 class ESBlock2(TheseusLayer):
-    def __init__(self,
-                 in_channels,
-                 out_channels):
+    def __init__(self, in_channels, out_channels):
         super().__init__()
 
         # branch1
@@ -205,9 +197,7 @@ class ESBlock2(TheseusLayer):
             kernel_size=3,
             groups=out_channels)
         self.concat_pw = ConvBNLayer(
-            in_channels=out_channels,
-            out_channels=out_channels,
-            kernel_size=1)
+            in_channels=out_channels, out_channels=out_channels, kernel_size=1)
 
     def forward(self, x):
         x1 = self.dw_1(x)
@@ -223,13 +213,20 @@ class ESBlock2(TheseusLayer):
 
 
 class ESNet(TheseusLayer):
-    def __init__(self, class_num=1000, scale=1.0, dropout_prob=0.2, class_expand=1280):
+    def __init__(self,
+                 class_num=1000,
+                 scale=1.0,
+                 dropout_prob=0.2,
+                 class_expand=1280):
         super().__init__()
         self.scale = scale
         self.class_num = class_num
         self.class_expand = class_expand
         stage_repeats = [3, 7, 3]
-        stage_out_channels = [-1, 24, make_divisible(116*scale), make_divisible(232*scale), make_divisible(464*scale), 1024]
+        stage_out_channels = [
+            -1, 24, make_divisible(116 * scale), make_divisible(232 * scale),
+            make_divisible(464 * scale), 1024
+        ]
 
         self.conv1 = ConvBNLayer(
             in_channels=3,
@@ -356,4 +353,3 @@ def ESNet_x1_0(pretrained=False, use_ssld=False, **kwargs):
     model = ESNet(scale=1.0, **kwargs)
     _load_pretrained(pretrained, model, MODEL_URLS["ESNet_x1_0"], use_ssld)
     return model
-