ta012/SSLAM_AS2M_Finetuned

1import torch2import torch.nn as nn3from timm.models.layers import trunc_normal_4from functools import partial5import numpy as np6from .model_core import (7    PatchEmbed_new,8    get_2d_sincos_pos_embed_flexible,9    FixedPositionalEncoder,10    AltBlock11)12 13class EAT(nn.Module):14    def __init__(self, config):15        super().__init__()16        self.config = config17        self.mode = config.model_variant  # "pretrain" or "finetune"18        19        # === Embedding / Encoder ===20        self.local_encoder = PatchEmbed_new(21            img_size=config.img_size,22            patch_size=config.patch_size,23            in_chans=config.in_chans,24            embed_dim=config.embed_dim,25            stride=config.stride26        )27 28        self.extra_tokens = nn.Parameter(torch.zeros(1, 1, config.embed_dim))29        self.pos_drop = nn.Dropout(p=config.drop_rate, inplace=True)30        trunc_normal_(self.extra_tokens, std=.02)31 32        self.fixed_positional_encoder = (33            FixedPositionalEncoder(self.build_sincos_pos_embed()) if config.fixed_positions else None34        )35 36        norm_layer = partial(nn.LayerNorm, eps=config.norm_eps, elementwise_affine=config.norm_affine)37        dpr = np.linspace(config.start_drop_path_rate, config.end_drop_path_rate, config.depth)38        self.blocks = nn.ModuleList([39            AltBlock(config.embed_dim, config.num_heads, config.mlp_ratio,40                     qkv_bias=config.qkv_bias, drop=config.drop_rate,41                     attn_drop=config.attn_drop_rate, mlp_drop=config.activation_dropout,42                     post_mlp_drop=config.post_mlp_drop, drop_path=dpr[i],43                     norm_layer=norm_layer, layer_norm_first=config.layer_norm_first,44                     ffn_targets=True)45            for i in range(config.depth)46        ])47 48        self.pre_norm = norm_layer(config.embed_dim)49 50        # === Head (for finetune) ===51        if self.mode == "finetune":52            self.fc_norm = nn.LayerNorm(config.embed_dim)53            self.head = nn.Linear(config.embed_dim, config.num_classes, bias=True)54        else:55            self.head = nn.Identity()56 57        self.apply(self._init_weights)58 59    def build_sincos_pos_embed(self):60        W = self.config.mel_bins // self.config.patch_size61        max_length = self.config.max_length62        embed_dim = self.config.embed_dim63        pos_embed = nn.Parameter(torch.zeros(1, max_length * W, embed_dim), requires_grad=False)64        emb = get_2d_sincos_pos_embed_flexible(embed_dim, (max_length, W), cls_token=False)65        pos_embed.data.copy_(torch.from_numpy(emb).float().unsqueeze(0))66        return pos_embed67 68    def _init_weights(self, m):69        if isinstance(m, nn.Linear):70            trunc_normal_(m.weight, std=.02)71            if m.bias is not None:72                nn.init.constant_(m.bias, 0)73        elif isinstance(m, nn.LayerNorm):74            nn.init.constant_(m.bias, 0)75            nn.init.constant_(m.weight, 1.0)76 77    def encode(self, x):78        B = x.shape[0]79        x = self.local_encoder(x)80        if self.fixed_positional_encoder is not None:81            x = x + self.fixed_positional_encoder(x, None)[:, :x.size(1), :]82        x = torch.cat((self.extra_tokens.expand(B, -1, -1), x), dim=1)83        x = self.pre_norm(x)84        x = self.pos_drop(x)85        for blk in self.blocks:86            x, _ = blk(x)87        return x88 89    def forward(self, x):90        x = self.encode(x)91        if self.mode == "finetune":92            x = x[:, 0]  # use cls token93            x = self.fc_norm(x)94            x = self.head(x)95        return x96 97    def extract_features(self, x):98        x = self.encode(x)99        return x100