Skip to content

van

model.backbone.van

PyTorch Visual Attention Network (VAN) model.

URL: https://github.com/Visual-Attention-Network/VAN-Classification

model_urls module-attribute 

model_urls = {'van_b0': 'https://huggingface.co/Visual-Attention-Network/VAN-Tiny-original/resolve/main/van_tiny_754.pth.tar', 'van_b1': 'https://huggingface.co/Visual-Attention-Network/VAN-Small-original/resolve/main/van_small_811.pth.tar', 'van_b2': 'https://huggingface.co/Visual-Attention-Network/VAN-Base-original/resolve/main/van_base_828.pth.tar', 'van_b3': 'https://huggingface.co/Visual-Attention-Network/VAN-Large-original/resolve/main/van_large_839.pth.tar'}

DWConv 

DWConv(dim=768)

Bases: Module

Source code in SaigeToolkit/model/backbone/van.py 
32
33
34
def __init__(self, dim=768):
    super(DWConv, self).__init__()
    self.dwconv = nn.Conv2d(dim, dim, 3, 1, 1, bias=True, groups=dim)

dwconv instance-attribute 

dwconv = Conv2d(dim, dim, 3, 1, 1, bias=True, groups=dim)

forward 

forward(x)
Source code in SaigeToolkit/model/backbone/van.py 
36
37
38
def forward(self, x):
    x = self.dwconv(x)
    return x

Mlp 

Mlp(in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.0)

Bases: Module

Source code in SaigeToolkit/model/backbone/van.py 
42
43
44
45
46
47
48
49
50
51
52
53
def __init__(
    self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.0
):
    super().__init__()
    out_features = out_features or in_features
    hidden_features = hidden_features or in_features
    self.fc1 = nn.Conv2d(in_features, hidden_features, 1)
    self.dwconv = DWConv(hidden_features)
    self.act = act_layer()
    self.fc2 = nn.Conv2d(hidden_features, out_features, 1)
    self.drop = nn.Dropout(drop)
    self.apply(self._init_weights)

fc1 instance-attribute 

fc1 = Conv2d(in_features, hidden_features, 1)

dwconv instance-attribute 

dwconv = DWConv(hidden_features)

act instance-attribute 

act = act_layer()

fc2 instance-attribute 

fc2 = Conv2d(hidden_features, out_features, 1)

drop instance-attribute 

drop = Dropout(drop)

_init_weights 

_init_weights(m)
Source code in SaigeToolkit/model/backbone/van.py 
55
56
57
58
59
60
61
62
63
64
65
66
67
68
def _init_weights(self, m):
    if isinstance(m, nn.Linear):
        trunc_normal_(m.weight, std=0.02)
        if isinstance(m, nn.Linear) and m.bias is not None:
            nn.init.constant_(m.bias, 0)
    elif isinstance(m, nn.LayerNorm):
        nn.init.constant_(m.bias, 0)
        nn.init.constant_(m.weight, 1.0)
    elif isinstance(m, nn.Conv2d):
        fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
        fan_out //= m.groups
        m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
        if m.bias is not None:
            m.bias.data.zero_()

forward 

forward(x)
Source code in SaigeToolkit/model/backbone/van.py 
70
71
72
73
74
75
76
77
def forward(self, x):
    x = self.fc1(x)
    x = self.dwconv(x)
    x = self.act(x)
    x = self.drop(x)
    x = self.fc2(x)
    x = self.drop(x)
    return x

LKA 

LKA(dim, clip_attn=None)

Bases: Module

Source code in SaigeToolkit/model/backbone/van.py 
81
82
83
84
85
86
def __init__(self, dim, clip_attn=None):
    super().__init__()
    self.conv0 = nn.Conv2d(dim, dim, 5, padding=2, groups=dim)
    self.conv_spatial = nn.Conv2d(dim, dim, 7, stride=1, padding=9, groups=dim, dilation=3)
    self.conv1 = nn.Conv2d(dim, dim, 1)
    self.clip_attn = clip_attn

conv0 instance-attribute 

conv0 = Conv2d(dim, dim, 5, padding=2, groups=dim)

conv_spatial instance-attribute 

conv_spatial = Conv2d(dim, dim, 7, stride=1, padding=9, groups=dim, dilation=3)

conv1 instance-attribute 

conv1 = Conv2d(dim, dim, 1)

clip_attn instance-attribute 

clip_attn = clip_attn

forward 

forward(x)
Source code in SaigeToolkit/model/backbone/van.py 
88
89
90
91
92
93
94
95
96
97
98
def forward(self, x):
    u = x.clone()
    attn = self.conv0(x)
    attn = self.conv_spatial(attn)
    attn = self.conv1(attn)
    attn = u * attn

    if self.clip_attn is not None:
        attn = torch.clamp(attn, min=-self.clip_attn, max=self.clip_attn)

    return attn

Attention 

Attention(d_model, clip_attn=None)

Bases: Module

Source code in SaigeToolkit/model/backbone/van.py 
102
103
104
105
106
107
108
def __init__(self, d_model, clip_attn=None):
    super().__init__()

    self.proj_1 = nn.Conv2d(d_model, d_model, 1)
    self.activation = nn.GELU()
    self.spatial_gating_unit = LKA(d_model, clip_attn=clip_attn)
    self.proj_2 = nn.Conv2d(d_model, d_model, 1)

proj_1 instance-attribute 

proj_1 = Conv2d(d_model, d_model, 1)

activation instance-attribute 

activation = GELU()

spatial_gating_unit instance-attribute 

spatial_gating_unit = LKA(d_model, clip_attn=clip_attn)

proj_2 instance-attribute 

proj_2 = Conv2d(d_model, d_model, 1)

forward 

forward(x)
Source code in SaigeToolkit/model/backbone/van.py 
110
111
112
113
114
115
116
117
def forward(self, x):
    shorcut = x.clone()
    x = self.proj_1(x)
    x = self.activation(x)
    x = self.spatial_gating_unit(x)
    x = self.proj_2(x)
    x = x + shorcut
    return x

Block 

Block(dim, mlp_ratio=4.0, drop=0.0, drop_path=0.0, act_layer=nn.GELU, clip_attn=None)

Bases: Module

Source code in SaigeToolkit/model/backbone/van.py 
121
122
123
124
125
126
127
128
129
130
131
132
133
134
def __init__(self, dim, mlp_ratio=4.0, drop=0.0, drop_path=0.0, act_layer=nn.GELU, clip_attn=None):
    super().__init__()
    self.norm1 = nn.BatchNorm2d(dim)
    self.attn = Attention(dim, clip_attn=clip_attn)
    self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()

    self.norm2 = nn.BatchNorm2d(dim)
    mlp_hidden_dim = int(dim * mlp_ratio)
    self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
    layer_scale_init_value = 1e-2
    self.layer_scale_1 = nn.Parameter(layer_scale_init_value * torch.ones((dim)), requires_grad=True)
    self.layer_scale_2 = nn.Parameter(layer_scale_init_value * torch.ones((dim)), requires_grad=True)

    self.apply(self._init_weights)

norm1 instance-attribute 

norm1 = BatchNorm2d(dim)

attn instance-attribute 

attn = Attention(dim, clip_attn=clip_attn)

drop_path instance-attribute 

drop_path = DropPath(drop_path) if drop_path > 0.0 else Identity()

norm2 instance-attribute 

norm2 = BatchNorm2d(dim)

mlp instance-attribute 

mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)

layer_scale_1 instance-attribute 

layer_scale_1 = Parameter(layer_scale_init_value * ones(dim), requires_grad=True)

layer_scale_2 instance-attribute 

layer_scale_2 = Parameter(layer_scale_init_value * ones(dim), requires_grad=True)

_init_weights 

_init_weights(m)
Source code in SaigeToolkit/model/backbone/van.py 
136
137
138
139
140
141
142
143
144
145
146
147
148
149
def _init_weights(self, m):
    if isinstance(m, nn.Linear):
        trunc_normal_(m.weight, std=0.02)
        if isinstance(m, nn.Linear) and m.bias is not None:
            nn.init.constant_(m.bias, 0)
    elif isinstance(m, nn.LayerNorm):
        nn.init.constant_(m.bias, 0)
        nn.init.constant_(m.weight, 1.0)
    elif isinstance(m, nn.Conv2d):
        fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
        fan_out //= m.groups
        m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
        if m.bias is not None:
            m.bias.data.zero_()

forward 

forward(x)
Source code in SaigeToolkit/model/backbone/van.py 
151
152
153
154
def forward(self, x):
    x = x + self.drop_path(self.layer_scale_1.unsqueeze(-1).unsqueeze(-1) * self.attn(self.norm1(x)))
    x = x + self.drop_path(self.layer_scale_2.unsqueeze(-1).unsqueeze(-1) * self.mlp(self.norm2(x)))
    return x

OverlapPatchEmbed 

OverlapPatchEmbed(img_size=224, patch_size=7, stride=4, in_channels=3, embed_dim=768)

Bases: Module

Image to Patch Embedding

Source code in SaigeToolkit/model/backbone/van.py 
160
161
162
163
164
165
166
167
168
169
170
171
172
def __init__(self, img_size=224, patch_size=7, stride=4, in_channels=3, embed_dim=768):
    super().__init__()
    patch_size = to_2tuple(patch_size)
    self.proj = nn.Conv2d(
        in_channels,
        embed_dim,
        kernel_size=patch_size,
        stride=stride,
        padding=(patch_size[0] // 2, patch_size[1] // 2),
    )
    self.norm = nn.BatchNorm2d(embed_dim)

    self.apply(self._init_weights)

proj instance-attribute 

proj = Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=stride, padding=(patch_size[0] // 2, patch_size[1] // 2))

norm instance-attribute 

norm = BatchNorm2d(embed_dim)

_init_weights 

_init_weights(m)
Source code in SaigeToolkit/model/backbone/van.py 
174
175
176
177
178
179
180
181
182
183
184
185
186
187
def _init_weights(self, m):
    if isinstance(m, nn.Linear):
        trunc_normal_(m.weight, std=0.02)
        if isinstance(m, nn.Linear) and m.bias is not None:
            nn.init.constant_(m.bias, 0)
    elif isinstance(m, nn.LayerNorm):
        nn.init.constant_(m.bias, 0)
        nn.init.constant_(m.weight, 1.0)
    elif isinstance(m, nn.Conv2d):
        fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
        fan_out //= m.groups
        m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
        if m.bias is not None:
            m.bias.data.zero_()

forward 

forward(x)
Source code in SaigeToolkit/model/backbone/van.py 
189
190
191
192
193
def forward(self, x):
    x = self.proj(x)
    _, _, H, W = x.shape
    x = self.norm(x)
    return x, H, W

VAN 

VAN(img_size: int = 224, in_channels: int = 3, embed_dims: List[int] = [64, 128, 256, 512], mlp_ratios: List[int] = [4, 4, 4, 4], drop_rate: float = 0.0, drop_path_rate: float = 0.0, norm_layer: Module = nn.LayerNorm, depths: List[int] = [3, 4, 6, 3], num_stages: int = 4, clip_attn: Optional[float] = None)

Bases: Module

Source code in SaigeToolkit/model/backbone/van.py 
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
def __init__(
    self,
    img_size: int = 224,
    in_channels: int = 3,
    embed_dims: List[int] = [64, 128, 256, 512],
    mlp_ratios: List[int] = [4, 4, 4, 4],
    drop_rate: float = 0.0,
    drop_path_rate: float = 0.0,
    norm_layer: nn.Module = nn.LayerNorm,
    depths: List[int] = [3, 4, 6, 3],
    num_stages: int = 4,
    clip_attn: Optional[float] = None,
):
    super().__init__()

    self.depths = depths
    self.num_stages = num_stages
    self.embed_dims = embed_dims

    dpr = [
        x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))
    ]  # stochastic depth decay rule
    cur = 0

    for i in range(num_stages):
        patch_embed = OverlapPatchEmbed(
            img_size=img_size if i == 0 else img_size // (2 ** (i + 1)),
            patch_size=7 if i == 0 else 3,
            stride=4 if i == 0 else 2,
            in_channels=in_channels if i == 0 else embed_dims[i - 1],
            embed_dim=embed_dims[i],
        )

        block = nn.ModuleList(
            [
                Block(
                    dim=embed_dims[i],
                    mlp_ratio=mlp_ratios[i],
                    drop=drop_rate,
                    drop_path=dpr[cur + j],
                    clip_attn=clip_attn,
                )
                for j in range(depths[i])
            ]
        )
        norm = norm_layer(embed_dims[i])
        cur += depths[i]

        setattr(self, f"patch_embed{i + 1}", patch_embed)
        setattr(self, f"block{i + 1}", block)
        setattr(self, f"norm{i + 1}", norm)

    self.apply(self._init_weights)

depths instance-attribute 

depths = depths

num_stages instance-attribute 

num_stages = num_stages

embed_dims instance-attribute 

embed_dims = embed_dims

_init_weights 

_init_weights(m)
Source code in SaigeToolkit/model/backbone/van.py 
251
252
253
254
255
256
257
258
259
260
261
262
263
264
def _init_weights(self, m):
    if isinstance(m, nn.Linear):
        trunc_normal_(m.weight, std=0.02)
        if isinstance(m, nn.Linear) and m.bias is not None:
            nn.init.constant_(m.bias, 0)
    elif isinstance(m, nn.LayerNorm):
        nn.init.constant_(m.bias, 0)
        nn.init.constant_(m.weight, 1.0)
    elif isinstance(m, nn.Conv2d):
        fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
        fan_out //= m.groups
        m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
        if m.bias is not None:
            m.bias.data.zero_()

freeze_patch_emb 

freeze_patch_emb()
Source code in SaigeToolkit/model/backbone/van.py 
266
267
def freeze_patch_emb(self):
    self.patch_embed1.requires_grad = False

no_weight_decay 

no_weight_decay()
Source code in SaigeToolkit/model/backbone/van.py 
269
270
271
272
273
274
275
276
277
@torch.jit.ignore
def no_weight_decay(self):
    return {
        "pos_embed1",
        "pos_embed2",
        "pos_embed3",
        "pos_embed4",
        "cls_token",
    }  # has pos_embed may be better

forward 

forward(x)
Source code in SaigeToolkit/model/backbone/van.py 
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
def forward(self, x):
    B = x.shape[0]

    for i in range(self.num_stages):
        patch_embed = getattr(self, f"patch_embed{i + 1}")
        block = getattr(self, f"block{i + 1}")
        norm = getattr(self, f"norm{i + 1}")
        x, H, W = patch_embed(x)
        for blk in block:
            x = blk(x)
        x = x.flatten(2).transpose(1, 2)
        x = norm(x)
        if i != self.num_stages - 1:
            x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()

    return x.mean(dim=1)

load_state_dict 

load_state_dict(state_dict, strict: bool = True)
Source code in SaigeToolkit/model/backbone/van.py 
296
297
298
299
300
301
def load_state_dict(self, state_dict, strict: bool = True):
    extend_state_dict_input_channel(state_dict, "patch_embed1.proj.weight", self.patch_embed1.proj)
    for key in list(state_dict):
        if key.startswith("head."):
            state_dict.pop(key)
    return super().load_state_dict(state_dict, strict)

extend_state_dict_input_channel 

extend_state_dict_input_channel(state_dict: Mapping[str, Any], input_weight_key: str, input_conv_layer: Conv2d) -> None

(Multipage) 3채널 이상인 이미지를 사용하기 위해 필요한 기능이며, state_dict의 input conv 채널이 네트워크의 input conv 채널보다 작은 경우 해당 weight의 채널을 늘려줍니다. 현재 네트워크가 가진 input conv 웨이트에서 앞 3 채널을 state_dict의 input conv 웨이트로 치환하는 방식을 사용합니다.

Parameters:

  • state_dict (Mapping[str, Any]) –

    로드하려는 weight

  • input_weight_key (str) –

    input conv weight의 이름

  • input_conv_layer (Conv2d) –

    input conv layer

Source code in SaigeToolkit/model/util.py 
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
def extend_state_dict_input_channel(
    state_dict: Mapping[str, Any],
    input_weight_key: str,
    input_conv_layer: nn.Conv2d,
) -> None:
    """(Multipage) 3채널 이상인 이미지를 사용하기 위해 필요한 기능이며,
    state_dict의 input conv 채널이 네트워크의 input conv 채널보다 작은 경우 해당 weight의 채널을 늘려줍니다.
    현재 네트워크가 가진 input conv 웨이트에서 앞 3 채널을 state_dict의 input conv 웨이트로 치환하는 방식을 사용합니다.

    Args:
        state_dict (Mapping[str, Any]): 로드하려는 weight
        input_weight_key (str): input conv weight의 이름
        input_conv_layer (nn.Conv2d): input conv layer
    """
    state_dict_in_channels = state_dict[input_weight_key].shape[1]
    if state_dict_in_channels != input_conv_layer.in_channels:
        state_dict_in_channels = min(state_dict_in_channels, input_conv_layer.in_channels)
        original_input_conv_weight = input_conv_layer.state_dict()["weight"]
        original_input_conv_weight[:, :state_dict_in_channels] = state_dict[input_weight_key][
            :, :state_dict_in_channels
        ]
        state_dict[input_weight_key] = original_input_conv_weight

load_model_weights 

load_model_weights(model, arch, kwargs)
Source code in SaigeToolkit/model/backbone/van.py 
312
313
314
315
316
317
318
319
320
321
def load_model_weights(model, arch, kwargs):
    url = model_urls[arch]
    checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu", check_hash=True)
    strict = True
    if "num_classes" in kwargs and kwargs["num_classes"] != 1000:
        strict = False
        del checkpoint["state_dict"]["head.weight"]
        del checkpoint["state_dict"]["head.bias"]
    model.load_state_dict(checkpoint["state_dict"], strict=strict)
    return model

van_b0 

van_b0(pretrained=False, **kwargs)
Source code in SaigeToolkit/model/backbone/van.py 
324
325
326
327
328
329
330
331
332
333
334
335
def van_b0(pretrained=False, **kwargs):
    model = VAN(
        embed_dims=[32, 64, 160, 256],
        mlp_ratios=[8, 8, 4, 4],
        norm_layer=partial(nn.LayerNorm, eps=1e-6),
        depths=[3, 3, 5, 2],
        **kwargs,
    )
    model.default_cfg = _cfg()
    if pretrained:
        model = load_model_weights(model, "van_b0", kwargs)
    return model

van_b1 

van_b1(pretrained=False, **kwargs)
Source code in SaigeToolkit/model/backbone/van.py 
338
339
340
341
342
343
344
345
346
347
348
349
def van_b1(pretrained=False, **kwargs):
    model = VAN(
        embed_dims=[64, 128, 320, 512],
        mlp_ratios=[8, 8, 4, 4],
        norm_layer=partial(nn.LayerNorm, eps=1e-6),
        depths=[2, 2, 4, 2],
        **kwargs,
    )
    model.default_cfg = _cfg()
    if pretrained:
        model = load_model_weights(model, "van_b1", kwargs)
    return model

van_b2 

van_b2(pretrained=False, **kwargs)
Source code in SaigeToolkit/model/backbone/van.py 
352
353
354
355
356
357
358
359
360
361
362
363
def van_b2(pretrained=False, **kwargs):
    model = VAN(
        embed_dims=[64, 128, 320, 512],
        mlp_ratios=[8, 8, 4, 4],
        norm_layer=partial(nn.LayerNorm, eps=1e-6),
        depths=[3, 3, 12, 3],
        **kwargs,
    )
    model.default_cfg = _cfg()
    if pretrained:
        model = load_model_weights(model, "van_b2", kwargs)
    return model

van_b3 

van_b3(pretrained=False, **kwargs)
Source code in SaigeToolkit/model/backbone/van.py 
366
367
368
369
370
371
372
373
374
375
376
377
def van_b3(pretrained=False, **kwargs):
    model = VAN(
        embed_dims=[64, 128, 320, 512],
        mlp_ratios=[8, 8, 4, 4],
        norm_layer=partial(nn.LayerNorm, eps=1e-6),
        depths=[3, 5, 27, 3],
        **kwargs,
    )
    model.default_cfg = _cfg()
    if pretrained:
        model = load_model_weights(model, "van_b3", kwargs)
    return model