| |
| import math |
|
|
| import torch |
| import torch.nn.functional as F |
| from torch import nn |
|
|
| from core.utils import get_init_fn |
| from core.vision_projector.base import BaseProjector |
|
|
|
|
| class AdaptiveAvgPooling(nn.Module): |
| def __init__(self, pooling_ratio=2): |
| super(AdaptiveAvgPooling, self).__init__() |
| self.pooling_ratio = pooling_ratio |
|
|
| def forward(self, x): |
| b, num_tokens, c = x.shape |
| h = int(math.sqrt(num_tokens)) |
| assert h * h == num_tokens |
|
|
| shape = (h // self.pooling_ratio, h // self.pooling_ratio) |
| x = x.permute(0, 2, 1).reshape(b, -1, h, h) |
| x = F.adaptive_avg_pool2d(x, shape) |
| x = x.flatten(2).transpose(1, 2) |
|
|
| return x |
|
|
|
|
| class MLPProjector(BaseProjector): |
| def __init__(self, args): |
| super().__init__() |
| self.setup_projector(args) |
| self.pooling_ratio = args.pooling_ratio |
| self.adaptive_avg_pool = AdaptiveAvgPooling(pooling_ratio=args.pooling_ratio) |
| self.remove_vision_class_token = args.remove_vision_class_token |
|
|
| def init_tensors(self): |
| self.init_method(self.projector[0].weight) |
| self.init_method(self.projector[0].bias) |
| self.init_method(self.projector[2].weight) |
| self.init_method(self.projector[2].bias) |
|
|
| def setup_projector(self, args): |
| self.init_method = get_init_fn(args.mlp_init, args.dim, init_depth=None) |
| input_size = args.vision_model["width"] |
| output_size = args.dim |
| self.projector = nn.Sequential( |
| nn.Linear( |
| in_features=input_size, |
| out_features=output_size, |
| bias=True, |
| dtype=torch.get_default_dtype(), |
| ), |
| nn.GELU(), |
| nn.Linear( |
| in_features=output_size, |
| out_features=output_size, |
| bias=True, |
| dtype=torch.get_default_dtype(), |
| ), |
| ) |
|
|
