|
|
|
|
|
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(), |
|
|
), |
|
|
) |
|
|
|
