distilled_vision_transformer.py 8.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272
  1. # copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. # Code was heavily based on https://github.com/facebookresearch/deit
  15. import paddle
  16. import paddle.nn as nn
  17. from .vision_transformer import VisionTransformer, Identity, trunc_normal_, zeros_
  18. from paddlex.ppcls.utils.save_load import load_dygraph_pretrain, load_dygraph_pretrain_from_url
  19. MODEL_URLS = {
  20. "DeiT_tiny_patch16_224":
  21. "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/DeiT_tiny_patch16_224_pretrained.pdparams",
  22. "DeiT_small_patch16_224":
  23. "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/DeiT_small_patch16_224_pretrained.pdparams",
  24. "DeiT_base_patch16_224":
  25. "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/DeiT_base_patch16_224_pretrained.pdparams",
  26. "DeiT_tiny_distilled_patch16_224":
  27. "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/DeiT_tiny_distilled_patch16_224_pretrained.pdparams",
  28. "DeiT_small_distilled_patch16_224":
  29. "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/DeiT_small_distilled_patch16_224_pretrained.pdparams",
  30. "DeiT_base_distilled_patch16_224":
  31. "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/DeiT_base_distilled_patch16_224_pretrained.pdparams",
  32. "DeiT_base_patch16_384":
  33. "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/DeiT_base_patch16_384_pretrained.pdparams",
  34. "DeiT_base_distilled_patch16_384":
  35. "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/DeiT_base_distilled_patch16_384_pretrained.pdparams",
  36. }
  37. __all__ = list(MODEL_URLS.keys())
  38. class DistilledVisionTransformer(VisionTransformer):
  39. def __init__(self,
  40. img_size=224,
  41. patch_size=16,
  42. class_num=1000,
  43. embed_dim=768,
  44. depth=12,
  45. num_heads=12,
  46. mlp_ratio=4,
  47. qkv_bias=False,
  48. norm_layer='nn.LayerNorm',
  49. epsilon=1e-5,
  50. **kwargs):
  51. super().__init__(
  52. img_size=img_size,
  53. patch_size=patch_size,
  54. class_num=class_num,
  55. embed_dim=embed_dim,
  56. depth=depth,
  57. num_heads=num_heads,
  58. mlp_ratio=mlp_ratio,
  59. qkv_bias=qkv_bias,
  60. norm_layer=norm_layer,
  61. epsilon=epsilon,
  62. **kwargs)
  63. self.pos_embed = self.create_parameter(
  64. shape=(1, self.patch_embed.num_patches + 2, self.embed_dim),
  65. default_initializer=zeros_)
  66. self.add_parameter("pos_embed", self.pos_embed)
  67. self.dist_token = self.create_parameter(
  68. shape=(1, 1, self.embed_dim), default_initializer=zeros_)
  69. self.add_parameter("cls_token", self.cls_token)
  70. self.head_dist = nn.Linear(
  71. self.embed_dim,
  72. self.class_num) if self.class_num > 0 else Identity()
  73. trunc_normal_(self.dist_token)
  74. trunc_normal_(self.pos_embed)
  75. self.head_dist.apply(self._init_weights)
  76. def forward_features(self, x):
  77. B = paddle.shape(x)[0]
  78. x = self.patch_embed(x)
  79. cls_tokens = self.cls_token.expand((B, -1, -1))
  80. dist_token = self.dist_token.expand((B, -1, -1))
  81. x = paddle.concat((cls_tokens, dist_token, x), axis=1)
  82. x = x + self.pos_embed
  83. x = self.pos_drop(x)
  84. for blk in self.blocks:
  85. x = blk(x)
  86. x = self.norm(x)
  87. return x[:, 0], x[:, 1]
  88. def forward(self, x):
  89. x, x_dist = self.forward_features(x)
  90. x = self.head(x)
  91. x_dist = self.head_dist(x_dist)
  92. return (x + x_dist) / 2
  93. def _load_pretrained(pretrained, model, model_url, use_ssld=False):
  94. if pretrained is False:
  95. pass
  96. elif pretrained is True:
  97. load_dygraph_pretrain_from_url(model, model_url, use_ssld=use_ssld)
  98. elif isinstance(pretrained, str):
  99. load_dygraph_pretrain(model, pretrained)
  100. else:
  101. raise RuntimeError(
  102. "pretrained type is not available. Please use `string` or `boolean` type."
  103. )
  104. def DeiT_tiny_patch16_224(pretrained=False, use_ssld=False, **kwargs):
  105. model = VisionTransformer(
  106. patch_size=16,
  107. embed_dim=192,
  108. depth=12,
  109. num_heads=3,
  110. mlp_ratio=4,
  111. qkv_bias=True,
  112. epsilon=1e-6,
  113. **kwargs)
  114. _load_pretrained(
  115. pretrained,
  116. model,
  117. MODEL_URLS["DeiT_tiny_patch16_224"],
  118. use_ssld=use_ssld)
  119. return model
  120. def DeiT_small_patch16_224(pretrained=False, use_ssld=False, **kwargs):
  121. model = VisionTransformer(
  122. patch_size=16,
  123. embed_dim=384,
  124. depth=12,
  125. num_heads=6,
  126. mlp_ratio=4,
  127. qkv_bias=True,
  128. epsilon=1e-6,
  129. **kwargs)
  130. _load_pretrained(
  131. pretrained,
  132. model,
  133. MODEL_URLS["DeiT_small_patch16_224"],
  134. use_ssld=use_ssld)
  135. return model
  136. def DeiT_base_patch16_224(pretrained=False, use_ssld=False, **kwargs):
  137. model = VisionTransformer(
  138. patch_size=16,
  139. embed_dim=768,
  140. depth=12,
  141. num_heads=12,
  142. mlp_ratio=4,
  143. qkv_bias=True,
  144. epsilon=1e-6,
  145. **kwargs)
  146. _load_pretrained(
  147. pretrained,
  148. model,
  149. MODEL_URLS["DeiT_base_patch16_224"],
  150. use_ssld=use_ssld)
  151. return model
  152. def DeiT_tiny_distilled_patch16_224(pretrained=False, use_ssld=False,
  153. **kwargs):
  154. model = DistilledVisionTransformer(
  155. patch_size=16,
  156. embed_dim=192,
  157. depth=12,
  158. num_heads=3,
  159. mlp_ratio=4,
  160. qkv_bias=True,
  161. epsilon=1e-6,
  162. **kwargs)
  163. _load_pretrained(
  164. pretrained,
  165. model,
  166. MODEL_URLS["DeiT_tiny_distilled_patch16_224"],
  167. use_ssld=use_ssld)
  168. return model
  169. def DeiT_small_distilled_patch16_224(pretrained=False,
  170. use_ssld=False,
  171. **kwargs):
  172. model = DistilledVisionTransformer(
  173. patch_size=16,
  174. embed_dim=384,
  175. depth=12,
  176. num_heads=6,
  177. mlp_ratio=4,
  178. qkv_bias=True,
  179. epsilon=1e-6,
  180. **kwargs)
  181. _load_pretrained(
  182. pretrained,
  183. model,
  184. MODEL_URLS["DeiT_small_distilled_patch16_224"],
  185. use_ssld=use_ssld)
  186. return model
  187. def DeiT_base_distilled_patch16_224(pretrained=False, use_ssld=False,
  188. **kwargs):
  189. model = DistilledVisionTransformer(
  190. patch_size=16,
  191. embed_dim=768,
  192. depth=12,
  193. num_heads=12,
  194. mlp_ratio=4,
  195. qkv_bias=True,
  196. epsilon=1e-6,
  197. **kwargs)
  198. _load_pretrained(
  199. pretrained,
  200. model,
  201. MODEL_URLS["DeiT_base_distilled_patch16_224"],
  202. use_ssld=use_ssld)
  203. return model
  204. def DeiT_base_patch16_384(pretrained=False, use_ssld=False, **kwargs):
  205. model = VisionTransformer(
  206. img_size=384,
  207. patch_size=16,
  208. embed_dim=768,
  209. depth=12,
  210. num_heads=12,
  211. mlp_ratio=4,
  212. qkv_bias=True,
  213. epsilon=1e-6,
  214. **kwargs)
  215. _load_pretrained(
  216. pretrained,
  217. model,
  218. MODEL_URLS["DeiT_base_patch16_384"],
  219. use_ssld=use_ssld)
  220. return model
  221. def DeiT_base_distilled_patch16_384(pretrained=False, use_ssld=False,
  222. **kwargs):
  223. model = DistilledVisionTransformer(
  224. img_size=384,
  225. patch_size=16,
  226. embed_dim=768,
  227. depth=12,
  228. num_heads=12,
  229. mlp_ratio=4,
  230. qkv_bias=True,
  231. epsilon=1e-6,
  232. **kwargs)
  233. _load_pretrained(
  234. pretrained,
  235. model,
  236. MODEL_URLS["DeiT_base_distilled_patch16_384"],
  237. use_ssld=use_ssld)
  238. return model