Merge pull request #1000 from cuicheng01/develop

修复levit初始化bug和tnt不能export model的问题

ppcls/arch/backbone/model_zoo/levit.py

@@ -80,7 +80,10 @@ class Linear_BN(nn.Sequential):
         super().__init__()
         self.add_sublayer('c', nn.Linear(a, b, bias_attr=False))
         bn = nn.BatchNorm1D(b)
-        ones_(bn.weight)
+        if bn_weight_init == 0:
+            zeros_(bn.weight)
+        else:
+            ones_(bn.weight)
         zeros_(bn.bias)
         self.add_sublayer('bn', bn)
 

ppcls/arch/backbone/model_zoo/tnt.py

@@ -44,7 +44,7 @@ def drop_path(x, drop_prob=0., training=False):
         return x
     keep_prob = paddle.to_tensor(1 - drop_prob)
     shape = (paddle.shape(x)[0], ) + (1, ) * (x.ndim - 1)
-    random_tensor = keep_prob + paddle.rand(shape, dtype=x.dtype)
+    random_tensor = paddle.add(keep_prob, paddle.rand(shape, dtype=x.dtype))
     random_tensor = paddle.floor(random_tensor)  # binarize
     output = x.divide(keep_prob) * random_tensor
     return output
@@ -114,14 +114,17 @@ class Attention(nn.Layer):
                 (2, 0, 3, 1, 4))
 
         q, k = qk[0], qk[1]
-        v = self.v(x).reshape((B, N, self.num_heads, -1)).transpose(
-            (0, 2, 1, 3))
+        v = self.v(x).reshape(
+            (B, N, self.num_heads, x.shape[-1] // self.num_heads)).transpose(
+                (0, 2, 1, 3))
 
-        attn = (q @k.transpose((0, 1, 3, 2))) * self.scale
+        attn = paddle.matmul(q, k.transpose((0, 1, 3, 2))) * self.scale
         attn = nn.functional.softmax(attn, axis=-1)
         attn = self.attn_drop(attn)
 
-        x = (attn @v).transpose((0, 2, 1, 3)).reshape((B, N, -1))
+        x = paddle.matmul(attn, v)
+        x = x.transpose((0, 2, 1, 3)).reshape(
+            (B, N, x.shape[-1] * x.shape[-3]))
         x = self.proj(x)
         x = self.proj_drop(x)
         return x
@@ -182,18 +185,22 @@ class Block(nn.Layer):
 
     def forward(self, pixel_embed, patch_embed):
         # inner
-        pixel_embed = pixel_embed + self.drop_path(
-            self.attn_in(self.norm_in(pixel_embed)))
-        pixel_embed = pixel_embed + self.drop_path(
-            self.mlp_in(self.norm_mlp_in(pixel_embed)))
+        pixel_embed = paddle.add(
+            pixel_embed,
+            self.drop_path(self.attn_in(self.norm_in(pixel_embed))))
+        pixel_embed = paddle.add(
+            pixel_embed,
+            self.drop_path(self.mlp_in(self.norm_mlp_in(pixel_embed))))
         # outer
         B, N, C = patch_embed.shape
-        patch_embed[:, 1:] = patch_embed[:, 1:] + self.proj(
-            self.norm1_proj(pixel_embed).reshape((B, N - 1, -1)))
-        patch_embed = patch_embed + self.drop_path(
-            self.attn_out(self.norm_out(patch_embed)))
-        patch_embed = patch_embed + self.drop_path(
-            self.mlp(self.norm_mlp(patch_embed)))
+        patch_embed[:, 1:] = paddle.add(
+            patch_embed[:, 1:],
+            self.proj(self.norm1_proj(pixel_embed).reshape((B, N - 1, -1))))
+        patch_embed = paddle.add(
+            patch_embed,
+            self.drop_path(self.attn_out(self.norm_out(patch_embed))))
+        patch_embed = paddle.add(
+            patch_embed, self.drop_path(self.mlp(self.norm_mlp(patch_embed))))
         return pixel_embed, patch_embed
 
 
@@ -222,11 +229,10 @@ class PixelEmbed(nn.Layer):
         x = self.proj(x)
         x = nn.functional.unfold(x, self.new_patch_size, self.new_patch_size)
         x = x.transpose((0, 2, 1)).reshape(
-            (B * self.num_patches, self.in_dim, self.new_patch_size,
-             self.new_patch_size))
+            (-1, self.in_dim, self.new_patch_size, self.new_patch_size))
         x = x + pixel_pos
-        x = x.reshape((B * self.num_patches, self.in_dim, -1)).transpose(
-            (0, 2, 1))
+        x = x.reshape((-1, self.in_dim, self.new_patch_size *
+                       self.new_patch_size)).transpose((0, 2, 1))
         return x
 
 
@@ -328,7 +334,8 @@ class TNT(nn.Layer):
         patch_embed = self.norm2_proj(
             self.proj(
                 self.norm1_proj(
-                    pixel_embed.reshape((B, self.num_patches, -1)))))
+                    pixel_embed.reshape((-1, self.num_patches, pixel_embed.
+                                         shape[-1] * pixel_embed.shape[-2])))))
         patch_embed = paddle.concat(
             (self.cls_token.expand((B, -1, -1)), patch_embed), axis=1)
         patch_embed = patch_embed + self.patch_pos