update

README.md

@@ -81,7 +81,7 @@ Note: Panoptic dataset (ADE, CityScapes, MV) are with the metric of PQ. Semantic
 
 | image encoder      | image size | VQAv2 | GQA  | VizWiz | SQA  | T-VQA | POPE | MME  | MM-Bench | MM-B-CN | SEED | LLaVA-Wild | MM-Vet |
 |---------------|----------|-------|------|--------|------|-------|------|------|----------|---------|------|------------|--------|
-| ViTamin-L   | 224 | 78.4  | 61.6 | 51.1   | 66.9 | 58.7  | 84.6 | 1421 | 65.4     | 58.4    | 57.7 | 64.5       | 33.6   |
+| ViTamin-L   | 336 | 78.4  | 61.6 | 51.1   | 66.9 | 58.7  | 84.6 | 1421 | 65.4     | 58.4    | 57.7 | 64.5       | 33.6   |
 | ViTamin-L   | 384 | 78.9  | 61.6 | 55.4   | 67.6 | 59.8  | 85.5 | 1447 | 64.5     | 58.3    | 57.9 | 66.1       | 33.6   |
 
 

vitamin.py

@@ -31,7 +31,7 @@ from torch.utils.checkpoint import checkpoint
 from timm.models.layers import create_attn, get_norm_layer, get_norm_act_layer, create_conv2d, make_divisible, trunc_normal_tf_
 
 
-from timm.layers import to_2tuple, DropPath, Format
+from timm.layers import to_2tuple, DropPath, Format # , trunc_normal_
 from timm.layers.norm_act import _create_act
 from timm.models._registry import register_model
 from timm.models._manipulate import named_apply, checkpoint_seq
@@ -335,6 +335,39 @@ class HybridEmbed(nn.Module):
         x = x.flatten(2).transpose(1, 2)
         return x
 
+def _trunc_normal_(tensor, mean, std, a, b):
+    # rewrite timm trunc normal
+    def norm_cdf(x):
+        # Computes standard normal cumulative distribution function
+        return (1. + math.erf(x / math.sqrt(2.))) / 2.
+
+    if (mean < a - 2 * std) or (mean > b + 2 * std):
+        warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
+                      "The distribution of values may be incorrect.",
+                      stacklevel=2)
+
+    l = norm_cdf((a - mean) / std)
+    u = norm_cdf((b - mean) / std)
+
+    # Uniformly fill tensor with values from [l, u], then translate to
+    # [2l-1, 2u-1].
+    tensor.uniform_(2 * l - 1, 2 * u - 1)
+
+    # Use inverse cdf transform for normal distribution to get truncated standard normal
+    # tensor.erfinv_() # NOTE: deleted as "erfinv_cuda" not implemented for 'BFloat16'
+
+    # Transform to proper mean, std
+    tensor.mul_(std * math.sqrt(2.))
+    tensor.add_(mean)
+
+    # Clamp to ensure it's in the proper range
+    tensor.clamp_(min=a, max=b)
+    return tensor
+
+def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
+    with torch.no_grad():
+        return _trunc_normal_(tensor, mean, std, a, b)
+
 class ViTamin(nn.Module):
     """ hack timm VisionTransformer
     """
@@ -563,6 +596,21 @@ class ViTamin(nn.Module):
             self.global_pool = global_pool
         self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()
 
+    def _pos_embed(self, x):
+        if self.no_embed_class:
+            # deit-3, updated JAX (big vision)
+            # position embedding does not overlap with class token, add then concat
+            x = x + self.pos_embed
+            if self.cls_token is not None:
+                x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1)
+        else:
+            # original timm, JAX, and deit vit impl
+            # pos_embed has entry for class token, concat then add
+            if self.cls_token is not None:
+                x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1)
+            x = x + self.pos_embed
+        return self.pos_drop(x)
+        
     def forward_features(self, x: torch.Tensor) -> torch.Tensor:
         x = self.patch_embed(x)
         if self.is_pos_embed: