model.py (14441B)
1 import torch 2 import torch.nn as nn 3 import csv 4 from torch.amp import autocast 5 import os 6 import wget 7 import timm 8 from timm.models.layers import to_2tuple, trunc_normal_ 9 10 os.environ["TORCH_HOME"] = "./pretrained_models" 11 12 13 # override the timm package to relax the input shape constraint. 14 class PatchEmbed(nn.Module): 15 def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768): 16 super().__init__() 17 18 img_size = to_2tuple(img_size) 19 patch_size = to_2tuple(patch_size) 20 num_patches = (img_size[1] // patch_size[1]) * ( 21 img_size[0] // patch_size[0] 22 ) 23 self.img_size = img_size 24 self.patch_size = patch_size 25 self.num_patches = num_patches 26 27 self.proj = nn.Conv2d( 28 in_chans, embed_dim, kernel_size=patch_size, stride=patch_size 29 ) 30 31 def forward(self, x): 32 x = self.proj(x).flatten(2).transpose(1, 2) 33 return x 34 35 36 class DistilledVisionTransformer: 37 def forward(self, x): 38 x, x_dist = self.forward_features(x) 39 x = self.head(x) 40 x_dist = self.head_dist(x_dist) 41 return (x + x_dist) / 2 42 43 44 class ASTModel(nn.Module): 45 """ 46 The AST model. 47 :param label_dim: the label dimension, i.e., the number of total classes, it is 527 for AudioSet, 50 for ESC-50, and 35 for speechcommands v2-35 48 :param fstride: the stride of patch spliting on the frequency dimension, for 16*16 patchs, fstride=16 means no overlap, fstride=10 means overlap of 6 49 :param tstride: the stride of patch spliting on the time dimension, for 16*16 patchs, tstride=16 means no overlap, tstride=10 means overlap of 6 50 :param input_fdim: the number of frequency bins of the input spectrogram 51 :param input_tdim: the number of time frames of the input spectrogram 52 :param imagenet_pretrain: if use ImageNet pretrained model 53 :param audioset_pretrain: if use full AudioSet and ImageNet pretrained model 54 :param model_size: the model size of AST, should be in [tiny224, small224, base224, base384], base224 and base 384 are same model, but are trained differently during ImageNet pretraining. 55 """ 56 57 def __init__( 58 self, 59 label_dim=527, 60 fstride=10, 61 tstride=10, 62 input_fdim=128, 63 input_tdim=1024, 64 imagenet_pretrain=False, 65 audioset_pretrain=False, 66 model_size="base384", 67 verbose=True, 68 ): 69 super(ASTModel, self).__init__() 70 assert ( 71 timm.__version__ == "0.4.5" 72 ), "Please use timm == 0.4.5, the code might not be compatible with newer versions." 73 74 if verbose == True: 75 print("---------------AST Model Summary---------------") 76 print( 77 "ImageNet pretraining: {:s}, AudioSet pretraining: {:s}".format( 78 str(imagenet_pretrain), str(audioset_pretrain) 79 ) 80 ) 81 # override timm input shape restriction 82 timm.models.vision_transformer.PatchEmbed = PatchEmbed 83 84 # fix torchscript compilation 85 timm.models.vision_transformer.DistilledVisionTransformer.forward = ( 86 DistilledVisionTransformer.forward 87 ) 88 89 # if AudioSet pretraining is not used (but ImageNet pretraining may still apply) 90 if audioset_pretrain == False: 91 if model_size == "tiny224": 92 self.v = timm.create_model( 93 "vit_deit_tiny_distilled_patch16_224", 94 pretrained=imagenet_pretrain, 95 ) 96 elif model_size == "small224": 97 self.v = timm.create_model( 98 "vit_deit_small_distilled_patch16_224", 99 pretrained=imagenet_pretrain, 100 ) 101 elif model_size == "base224": 102 self.v = timm.create_model( 103 "vit_deit_base_distilled_patch16_224", 104 pretrained=imagenet_pretrain, 105 ) 106 elif model_size == "base384": 107 self.v = timm.create_model( 108 "vit_deit_base_distilled_patch16_384", 109 pretrained=imagenet_pretrain, 110 ) 111 else: 112 raise Exception( 113 "Model size must be one of tiny224, small224, base224, base384." 114 ) 115 self.original_num_patches = self.v.patch_embed.num_patches 116 self.oringal_hw = int(self.original_num_patches**0.5) 117 self.original_embedding_dim = self.v.pos_embed.shape[2] 118 self.mlp_head = nn.Sequential( 119 nn.LayerNorm(self.original_embedding_dim), 120 nn.Linear(self.original_embedding_dim, label_dim), 121 ) 122 123 # automatcially get the intermediate shape 124 f_dim, t_dim = self.get_shape( 125 fstride, tstride, input_fdim, input_tdim 126 ) 127 num_patches = f_dim * t_dim 128 self.v.patch_embed.num_patches = num_patches 129 if verbose == True: 130 print( 131 "frequncey stride={:d}, time stride={:d}".format( 132 fstride, tstride 133 ) 134 ) 135 print("number of patches={:d}".format(num_patches)) 136 137 # the linear projection layer 138 new_proj = torch.nn.Conv2d( 139 1, 140 self.original_embedding_dim, 141 kernel_size=(16, 16), 142 stride=(fstride, tstride), 143 ) 144 if imagenet_pretrain == True: 145 new_proj.weight = torch.nn.Parameter( 146 torch.sum(self.v.patch_embed.proj.weight, dim=1).unsqueeze( 147 1 148 ) 149 ) 150 new_proj.bias = self.v.patch_embed.proj.bias 151 self.v.patch_embed.proj = new_proj 152 153 # the positional embedding 154 if imagenet_pretrain == True: 155 # get the positional embedding from deit model, skip the first two tokens (cls token and distillation token), reshape it to original 2D shape (24*24). 156 new_pos_embed = ( 157 self.v.pos_embed[:, 2:, :] 158 .detach() 159 .reshape( 160 1, 161 self.original_num_patches, 162 self.original_embedding_dim, 163 ) 164 .transpose(1, 2) 165 .reshape( 166 1, 167 self.original_embedding_dim, 168 self.oringal_hw, 169 self.oringal_hw, 170 ) 171 ) 172 # cut (from middle) or interpolate the second dimension of the positional embedding 173 if t_dim <= self.oringal_hw: 174 new_pos_embed = new_pos_embed[ 175 :, 176 :, 177 :, 178 int(self.oringal_hw / 2) 179 - int(t_dim / 2) : int(self.oringal_hw / 2) 180 - int(t_dim / 2) 181 + t_dim, 182 ] 183 else: 184 new_pos_embed = torch.nn.functional.interpolate( 185 new_pos_embed, 186 size=(self.oringal_hw, t_dim), 187 mode="bilinear", 188 ) 189 # cut (from middle) or interpolate the first dimension of the positional embedding 190 if f_dim <= self.oringal_hw: 191 new_pos_embed = new_pos_embed[ 192 :, 193 :, 194 int(self.oringal_hw / 2) 195 - int(f_dim / 2) : int(self.oringal_hw / 2) 196 - int(f_dim / 2) 197 + f_dim, 198 :, 199 ] 200 else: 201 new_pos_embed = torch.nn.functional.interpolate( 202 new_pos_embed, size=(f_dim, t_dim), mode="bilinear" 203 ) 204 # flatten the positional embedding 205 new_pos_embed = new_pos_embed.reshape( 206 1, self.original_embedding_dim, num_patches 207 ).transpose(1, 2) 208 # concatenate the above positional embedding with the cls token and distillation token of the deit model. 209 self.v.pos_embed = nn.Parameter( 210 torch.cat( 211 [self.v.pos_embed[:, :2, :].detach(), new_pos_embed], 212 dim=1, 213 ) 214 ) 215 else: 216 # if not use imagenet pretrained model, just randomly initialize a learnable positional embedding 217 # TODO can use sinusoidal positional embedding instead 218 new_pos_embed = nn.Parameter( 219 torch.zeros( 220 1, 221 self.v.patch_embed.num_patches + 2, 222 self.original_embedding_dim, 223 ) 224 ) 225 self.v.pos_embed = new_pos_embed 226 trunc_normal_(self.v.pos_embed, std=0.02) 227 228 # now load a model that is pretrained on both ImageNet and AudioSet 229 elif audioset_pretrain == True: 230 if audioset_pretrain == True and imagenet_pretrain == False: 231 raise ValueError( 232 "currently model pretrained on only audioset is not supported, please set imagenet_pretrain = True to use audioset pretrained model." 233 ) 234 if model_size != "base384": 235 raise ValueError( 236 "currently only has base384 AudioSet pretrained model." 237 ) 238 device = torch.device( 239 "cuda" if torch.cuda.is_available() else "cpu" 240 ) 241 if ( 242 os.path.exists("./pretrained_models/audioset_10_10_0.4593.pth") 243 == False 244 ): 245 # this model performs 0.4593 mAP on the audioset eval set 246 audioset_mdl_url = "https://www.dropbox.com/s/cv4knew8mvbrnvq/audioset_0.4593.pth?dl=1" 247 wget.download( 248 audioset_mdl_url, 249 out="./pretrained_models/audioset_10_10_0.4593.pth", 250 ) 251 sd = torch.load( 252 "./pretrained_models/audioset_10_10_0.4593.pth", 253 map_location=device, 254 ) 255 audio_model = ASTModel( 256 label_dim=527, 257 fstride=10, 258 tstride=10, 259 input_fdim=128, 260 input_tdim=1024, 261 imagenet_pretrain=False, 262 audioset_pretrain=False, 263 model_size="base384", 264 verbose=False, 265 ) 266 audio_model = torch.nn.DataParallel(audio_model) 267 audio_model.load_state_dict(sd, strict=False) 268 self.v = audio_model.module.v 269 self.original_embedding_dim = self.v.pos_embed.shape[2] 270 self.mlp_head = nn.Sequential( 271 nn.LayerNorm(self.original_embedding_dim), 272 nn.Linear(self.original_embedding_dim, label_dim), 273 ) 274 275 f_dim, t_dim = self.get_shape( 276 fstride, tstride, input_fdim, input_tdim 277 ) 278 num_patches = f_dim * t_dim 279 self.v.patch_embed.num_patches = num_patches 280 if verbose == True: 281 print( 282 "frequncey stride={:d}, time stride={:d}".format( 283 fstride, tstride 284 ) 285 ) 286 print("number of patches={:d}".format(num_patches)) 287 288 new_pos_embed = ( 289 self.v.pos_embed[:, 2:, :] 290 .detach() 291 .reshape(1, 1212, 768) 292 .transpose(1, 2) 293 .reshape(1, 768, 12, 101) 294 ) 295 # if the input sequence length is larger than the original audioset (10s), then cut the positional embedding 296 if t_dim < 101: 297 new_pos_embed = new_pos_embed[ 298 :, :, :, 50 - int(t_dim / 2) : 50 - int(t_dim / 2) + t_dim 299 ] 300 # otherwise interpolate 301 else: 302 new_pos_embed = torch.nn.functional.interpolate( 303 new_pos_embed, size=(12, t_dim), mode="bilinear" 304 ) 305 if f_dim < 12: 306 new_pos_embed = new_pos_embed[ 307 :, :, 6 - int(f_dim / 2) : 6 - int(f_dim / 2) + f_dim, : 308 ] 309 # otherwise interpolate 310 elif f_dim > 12: 311 new_pos_embed = torch.nn.functional.interpolate( 312 new_pos_embed, size=(f_dim, t_dim), mode="bilinear" 313 ) 314 new_pos_embed = new_pos_embed.reshape( 315 1, 768, num_patches 316 ).transpose(1, 2) 317 self.v.pos_embed = nn.Parameter( 318 torch.cat( 319 [self.v.pos_embed[:, :2, :].detach(), new_pos_embed], dim=1 320 ) 321 ) 322 323 def get_shape(self, fstride, tstride, input_fdim=128, input_tdim=1024): 324 test_input = torch.randn(1, 1, input_fdim, input_tdim) 325 test_proj = nn.Conv2d( 326 1, 327 self.original_embedding_dim, 328 kernel_size=(16, 16), 329 stride=(fstride, tstride), 330 ) 331 test_out = test_proj(test_input) 332 f_dim = test_out.shape[2] 333 t_dim = test_out.shape[3] 334 return f_dim, t_dim 335 336 @autocast("cuda") 337 def forward(self, x): 338 """ 339 :param x: the input spectrogram, expected shape: (batch_size, time_frame_num, frequency_bins), e.g., (12, 1024, 128) 340 :return: prediction 341 """ 342 # expect input x = (batch_size, time_frame_num, frequency_bins), e.g., (12, 1024, 128) 343 x = x.unsqueeze(1) 344 x = x.transpose(2, 3) 345 346 B = x.shape[0] 347 x = self.v.patch_embed(x) 348 cls_tokens = self.v.cls_token.expand(B, -1, -1) 349 dist_token = self.v.dist_token.expand(B, -1, -1) 350 x = torch.cat((cls_tokens, dist_token, x), dim=1) 351 x = x + self.v.pos_embed 352 x = self.v.pos_drop(x) 353 for blk in self.v.blocks: 354 x = blk(x) 355 x = self.v.norm(x) 356 x = (x[:, 0] + x[:, 1]) / 2 357 358 x = self.mlp_head(x) 359 360 return x