Create mobilenet_v1.py

ppcls/arch/backbone/legendary_models/mobilenet_v1.py

@@ -18,7 +18,7 @@ import numpy as np
 import paddle
 from paddle import ParamAttr
 import paddle.nn as nn
-from paddle.nn import Conv2D, BatchNorm, Linear, Dropout, ReLU
+from paddle.nn import Conv2D, BatchNorm, Linear, Dropout, ReLU, Flatten
 from paddle.nn import AdaptiveAvgPool2D, MaxPool2D, AvgPool2D
 from paddle.nn.initializer import KaimingNormal
 import math
@@ -131,17 +131,20 @@ class MobileNet(TheseusLayer):
                             scale=scale) for params in self.cfg])
 
         self.pool2d_avg = AdaptiveAvgPool2D(1)
+        self.flatten    = Flatten(start_axis=1, stop_axis=-1)
 
         self.out = Linear(
             int(1024 * scale),
             class_dim,
             weight_attr=ParamAttr(initializer=KaimingNormal()))
+        
+        
 
     def forward(self, x):
         x = self.conv1(x)
         x = self.blocks(x)
         x = self.pool2d_avg(x)
-        x = paddle.flatten(x, start_axis=1, stop_axis=-1)
+        x = self.flatten(x)
         x = self.out(x)
         return x