Update distilled_vision_transformer.py

ppcls/modeling/architectures/distilled_vision_transformer.py

@@ -16,7 +16,6 @@ import paddle
 import paddle.nn as nn
 from .vision_transformer import VisionTransformer, Identity, trunc_normal_, zeros_
 
-
 __all__ = [
     'DeiT_tiny_patch16_224', 'DeiT_small_patch16_224', 'DeiT_base_patch16_224',
     'DeiT_tiny_distilled_patch16_224', 'DeiT_small_distilled_patch16_224',
@@ -26,14 +25,33 @@ __all__ = [
 
 
 class DistilledVisionTransformer(VisionTransformer):
-    def __init__(self, img_size=224, patch_size=16, class_dim=1000, embed_dim=768, depth=12,
-                 num_heads=12, mlp_ratio=4, qkv_bias=False, norm_layer='nn.LayerNorm', epsilon=1e-5,
+    def __init__(self,
+                 img_size=224,
+                 patch_size=16,
+                 class_dim=1000,
+                 embed_dim=768,
+                 depth=12,
+                 num_heads=12,
+                 mlp_ratio=4,
+                 qkv_bias=False,
+                 norm_layer='nn.LayerNorm',
+                 epsilon=1e-5,
                  **kwargs):
-        super().__init__(img_size=img_size, patch_size=patch_size, class_dim=class_dim, embed_dim=embed_dim, depth=depth,
-                         num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, norm_layer=norm_layer, epsilon=epsilon,
-                         **kwargs)
+        super().__init__(
+            img_size=img_size,
+            patch_size=patch_size,
+            class_dim=class_dim,
+            embed_dim=embed_dim,
+            depth=depth,
+            num_heads=num_heads,
+            mlp_ratio=mlp_ratio,
+            qkv_bias=qkv_bias,
+            norm_layer=norm_layer,
+            epsilon=epsilon,
+            **kwargs)
         self.pos_embed = self.create_parameter(
-            shape=(1, self.patch_embed.num_patches + 2, self.embed_dim), default_initializer=zeros_)
+            shape=(1, self.patch_embed.num_patches + 2, self.embed_dim),
+            default_initializer=zeros_)
         self.add_parameter("pos_embed", self.pos_embed)
 
         self.dist_token = self.create_parameter(
@@ -41,14 +59,15 @@ class DistilledVisionTransformer(VisionTransformer):
         self.add_parameter("cls_token", self.cls_token)
 
         self.head_dist = nn.Linear(
-            self.embed_dim, self.class_dim) if self.class_dim > 0 else Identity()
+            self.embed_dim,
+            self.class_dim) if self.class_dim > 0 else Identity()
 
         trunc_normal_(self.dist_token)
         trunc_normal_(self.pos_embed)
         self.head_dist.apply(self._init_weights)
 
     def forward_features(self, x):
-        B = x.shape[0]
+        B = paddle.shape(x)[0]
         x = self.patch_embed(x)
 
         cls_tokens = self.cls_token.expand((B, -1, -1))
@@ -73,55 +92,105 @@ class DistilledVisionTransformer(VisionTransformer):
 
 def DeiT_tiny_patch16_224(**kwargs):
     model = VisionTransformer(
-        patch_size=16, embed_dim=192, depth=12, num_heads=3, mlp_ratio=4, qkv_bias=True,
-        epsilon=1e-6, **kwargs)
+        patch_size=16,
+        embed_dim=192,
+        depth=12,
+        num_heads=3,
+        mlp_ratio=4,
+        qkv_bias=True,
+        epsilon=1e-6,
+        **kwargs)
     return model
 
 
 def DeiT_small_patch16_224(**kwargs):
     model = VisionTransformer(
-        patch_size=16, embed_dim=384, depth=12, num_heads=6, mlp_ratio=4, qkv_bias=True,
-        epsilon=1e-6, **kwargs)
+        patch_size=16,
+        embed_dim=384,
+        depth=12,
+        num_heads=6,
+        mlp_ratio=4,
+        qkv_bias=True,
+        epsilon=1e-6,
+        **kwargs)
     return model
 
 
 def DeiT_base_patch16_224(**kwargs):
     model = VisionTransformer(
-        patch_size=16, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4, qkv_bias=True,
-        epsilon=1e-6, **kwargs)
+        patch_size=16,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4,
+        qkv_bias=True,
+        epsilon=1e-6,
+        **kwargs)
     return model
 
 
 def DeiT_tiny_distilled_patch16_224(**kwargs):
     model = DistilledVisionTransformer(
-        patch_size=16, embed_dim=192, depth=12, num_heads=3, mlp_ratio=4, qkv_bias=True,
-        epsilon=1e-6, **kwargs)
+        patch_size=16,
+        embed_dim=192,
+        depth=12,
+        num_heads=3,
+        mlp_ratio=4,
+        qkv_bias=True,
+        epsilon=1e-6,
+        **kwargs)
     return model
 
 
 def DeiT_small_distilled_patch16_224(**kwargs):
     model = DistilledVisionTransformer(
-        patch_size=16, embed_dim=384, depth=12, num_heads=6, mlp_ratio=4, qkv_bias=True,
-        epsilon=1e-6, **kwargs)
+        patch_size=16,
+        embed_dim=384,
+        depth=12,
+        num_heads=6,
+        mlp_ratio=4,
+        qkv_bias=True,
+        epsilon=1e-6,
+        **kwargs)
     return model
 
 
 def DeiT_base_distilled_patch16_224(**kwargs):
     model = DistilledVisionTransformer(
-        patch_size=16, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4, qkv_bias=True,
-        epsilon=1e-6, **kwargs)
+        patch_size=16,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4,
+        qkv_bias=True,
+        epsilon=1e-6,
+        **kwargs)
     return model
 
 
 def DeiT_base_patch16_384(**kwargs):
     model = VisionTransformer(
-        img_size=384, patch_size=16, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4, qkv_bias=True,
-        epsilon=1e-6, **kwargs)
+        img_size=384,
+        patch_size=16,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4,
+        qkv_bias=True,
+        epsilon=1e-6,
+        **kwargs)
     return model
 
 
 def DeiT_base_distilled_patch16_384(**kwargs):
     model = DistilledVisionTransformer(
-        img_size=384, patch_size=16, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4, qkv_bias=True,
-        epsilon=1e-6, **kwargs)
+        img_size=384,
+        patch_size=16,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4,
+        qkv_bias=True,
+        epsilon=1e-6,
+        **kwargs)
     return model