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
def __init__(self, dim=768):
    super(DWConv, self).__init__()
    self.dwconv = nn.Conv2d(dim, dim, 3, 1, 1, bias=True, groups=dim)

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
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)

LKA

LKA(dim, clip_attn=None)

Bases: Module

Source code in SaigeToolkit/model/backbone/van.py
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

Attention

Attention(d_model, clip_attn=None)

Bases: Module

Source code in SaigeToolkit/model/backbone/van.py
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)

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
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)

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
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)

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
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)

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
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
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
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
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
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
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