|
|
import torch
|
|
|
import torch.nn as nn
|
|
|
from clip.model import Transformer
|
|
|
from mmengine.model import BaseModule
|
|
|
|
|
|
|
|
|
class TransformerAdapter(BaseModule):
|
|
|
|
|
|
def __init__(self,
|
|
|
clip_model: nn.Module,
|
|
|
num_segs: int,
|
|
|
num_layers: int = 6):
|
|
|
super(TransformerAdapter, self).__init__()
|
|
|
self.num_segs = num_segs
|
|
|
|
|
|
embed_dim = clip_model.text_projection.shape[1]
|
|
|
transformer_width = clip_model.ln_final.weight.shape[0]
|
|
|
transformer_heads = transformer_width // 64
|
|
|
|
|
|
self.frame_position_embeddings = nn.Embedding(self.num_segs, embed_dim)
|
|
|
self.transformer = Transformer(
|
|
|
width=embed_dim, layers=num_layers, heads=transformer_heads)
|
|
|
|
|
|
def init_weights(self):
|
|
|
for module in self.modules():
|
|
|
if isinstance(module, (nn.Linear, nn.Embedding)):
|
|
|
module.weight.data.normal_(mean=0.0, std=0.02)
|
|
|
elif isinstance(module, nn.LayerNorm):
|
|
|
module.bias.data.zero_()
|
|
|
module.weight.data.fill_(1.0)
|
|
|
if isinstance(module, nn.Linear) and module.bias is not None:
|
|
|
module.bias.data.zero_()
|
|
|
|
|
|
def forward(self, x: torch.Tensor):
|
|
|
b, seq_length, c = x.size()
|
|
|
|
|
|
x_original = x
|
|
|
position_ids = torch.arange(
|
|
|
seq_length, dtype=torch.long, device=x.device)
|
|
|
embeddings = self.frame_position_embeddings(position_ids)
|
|
|
x = x + embeddings.unsqueeze(0)
|
|
|
x = x.transpose(0, 1)
|
|
|
x = self.transformer(x)
|
|
|
x = x.transpose(0, 1)
|
|
|
x = x.type(x_original.dtype) + x_original
|
|
|
return x.mean(dim=1)
|
|
|
|
