rebuild

model/vit.py

@@ -0,0 +1,42 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+import torch
+import torch.nn as nn
+from functools import partial, reduce
+from operator import mul
+
+from timm.models.vision_transformer import VisionTransformer, _cfg
+
+__all__ = [
+    'vit_small',
+    'vit_base',
+]
+
+
+def vit_small(**kwargs):
+    model = VisionTransformer(
+        patch_size=16, embed_dim=384, depth=12, num_heads=12, mlp_ratio=4, qkv_bias=True, num_classes=256,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    # model.default_cfg = _cfg()
+    return model
+
+
+def vit_base(**kwargs):
+    model = VisionTransformer(
+        patch_size=16, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4, qkv_bias=True, num_classes=256,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    model.default_cfg = _cfg(num_classes=256)
+    return model
+
+
+if __name__ == '__main__':
+    img = torch.randn(8, 3, 224, 224)
+    vit = vit_base()
+    out = vit(img)
+    print(out.shape)
+    # print(count_parameters(vit))