fix a bug for using smaller fstride with pretrained model

YuanGongND · Jul 9, 2022 · d7d8b4b · d7d8b4b
1 parent 87a8004
commit d7d8b4b
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/src/models/ast_models.py b/src/models/ast_models.py
@@ -145,6 +145,11 @@ def __init__(self, label_dim=527, fstride=10, tstride=10, input_fdim=128, input_
             # otherwise interpolate
             else:
                 new_pos_embed = torch.nn.functional.interpolate(new_pos_embed, size=(12, t_dim), mode='bilinear')
+            if f_dim < 12:
+                new_pos_embed = new_pos_embed[:, :, 6 - int(f_dim/2): 6 - int(f_dim/2) + f_dim, :]
+            # otherwise interpolate
+            elif f_dim > 12:
+                new_pos_embed = torch.nn.functional.interpolate(new_pos_embed, size=(f_dim, t_dim), mode='bilinear')
             new_pos_embed = new_pos_embed.reshape(1, 768, num_patches).transpose(1, 2)
             self.v.pos_embed = nn.Parameter(torch.cat([self.v.pos_embed[:, :2, :].detach(), new_pos_embed], dim=1))