forked from pengzhiliang/Conformer
-
Notifications
You must be signed in to change notification settings - Fork 0
/
models.py
110 lines (97 loc) · 4.04 KB
/
models.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
import torch
import torch.nn as nn
from functools import partial
# from timm.models.vision_transformer import VisionTransformer, _cfg
from vision_transformer import VisionTransformer, _cfg
from conformer import Conformer
from timm.models.registry import register_model
@register_model
def deit_tiny_patch16_224(pretrained=False, **kwargs):
model = VisionTransformer(
patch_size=16, embed_dim=192, depth=12, num_heads=3, mlp_ratio=4, qkv_bias=True,
norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
model.default_cfg = _cfg()
if pretrained:
checkpoint = torch.hub.load_state_dict_from_url(
url="https://dl.fbaipublicfiles.com/deit/deit_tiny_patch16_224-a1311bcf.pth",
map_location="cpu", check_hash=True
)
model.load_state_dict(checkpoint["model"])
return model
@register_model
def deit_small_patch16_224(pretrained=False, **kwargs):
model = VisionTransformer(
patch_size=16, embed_dim=384, depth=12, num_heads=6, mlp_ratio=4, qkv_bias=True,
norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
model.default_cfg = _cfg()
if pretrained:
checkpoint = torch.hub.load_state_dict_from_url(
url="https://dl.fbaipublicfiles.com/deit/deit_small_patch16_224-cd65a155.pth",
map_location="cpu", check_hash=True
)
model.load_state_dict(checkpoint["model"])
return model
@register_model
def deit_med_patch16_224(pretrained=False, **kwargs):
model = VisionTransformer(
patch_size=16, embed_dim=576, depth=12, num_heads=9, mlp_ratio=4, qkv_bias=True,
norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
model.default_cfg = _cfg()
if pretrained:
raise NotImplementedError
return model
@register_model
def deit_base_patch16_224(pretrained=False, **kwargs):
model = VisionTransformer(
patch_size=16, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4, qkv_bias=True,
norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
model.default_cfg = _cfg()
if pretrained:
checkpoint = torch.hub.load_state_dict_from_url(
url="https://dl.fbaipublicfiles.com/deit/deit_base_patch16_224-b5f2ef4d.pth",
map_location="cpu", check_hash=True
)
model.load_state_dict(checkpoint["model"])
return model
@register_model
def Conformer_tinytinytiny_patch16(pretrained=False, **kwargs):
model = Conformer(patch_size=16, channel_ratio=1, embed_dim=192, depth=3,
num_heads=3, mlp_ratio=4, qkv_bias=True, **kwargs)
if pretrained:
raise NotImplementedError
return model
@register_model
def Conformer_tinytiny_patch16(pretrained=False, **kwargs):
model = Conformer(patch_size=16, channel_ratio=1, embed_dim=384, depth=6,
num_heads=3, mlp_ratio=4, qkv_bias=True, **kwargs)
if pretrained:
raise NotImplementedError
return model
@register_model
def Conformer_tiny_patch16(pretrained=False, **kwargs):
model = Conformer(patch_size=16, channel_ratio=1, embed_dim=384, depth=12,
num_heads=6, mlp_ratio=4, qkv_bias=True, **kwargs)
if pretrained:
raise NotImplementedError
return model
@register_model
def Conformer_small_patch16(pretrained=False, **kwargs):
model = Conformer(patch_size=16, channel_ratio=4, embed_dim=384, depth=12,
num_heads=6, mlp_ratio=4, qkv_bias=True, **kwargs)
if pretrained:
raise NotImplementedError
return model
@register_model
def Conformer_small_patch32(pretrained=False, **kwargs):
model = Conformer(patch_size=32, channel_ratio=4, embed_dim=384, depth=12,
num_heads=6, mlp_ratio=4, qkv_bias=True, **kwargs)
if pretrained:
raise NotImplementedError
return model
@register_model
def Conformer_base_patch16(pretrained=False, **kwargs):
model = Conformer(patch_size=16, channel_ratio=6, embed_dim=576, depth=12,
num_heads=9, mlp_ratio=4, qkv_bias=True, **kwargs)
if pretrained:
raise NotImplementedError
return model