56 lines
1.7 KiB
Python
56 lines
1.7 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
class PostfuseModule(nn.Module):
|
|
def __init__(self, embed_dim, embed_dim_img):
|
|
super().__init__()
|
|
self.mlp1 = MLP(embed_dim_img, embed_dim, embed_dim, use_residual=False)
|
|
self.mlp2 = MLP(embed_dim, embed_dim, embed_dim, use_residual=True)
|
|
self.layer_norm = nn.LayerNorm(embed_dim)
|
|
|
|
@property
|
|
def dtype(self):
|
|
try:
|
|
return next(self.parameters()).dtype
|
|
except StopIteration:
|
|
return torch.float32
|
|
|
|
def fuse_fn(self, object_embeds):
|
|
text_object_embeds = self.mlp1(object_embeds)
|
|
text_object_embeds = self.mlp2(text_object_embeds)
|
|
text_object_embeds = self.layer_norm(text_object_embeds)
|
|
return text_object_embeds
|
|
|
|
def forward(
|
|
self,
|
|
text_embeds,
|
|
object_embeds,
|
|
fuse_index,
|
|
) -> torch.Tensor:
|
|
text_object_embed = self.fuse_fn(object_embeds)
|
|
text_embeds_new = text_embeds.clone()
|
|
text_embeds_new[:, fuse_index, :] = text_object_embed.squeeze(1)
|
|
|
|
return text_embeds_new
|
|
|
|
class MLP(nn.Module):
|
|
def __init__(self, in_dim, out_dim, hidden_dim, use_residual=True):
|
|
super().__init__()
|
|
if use_residual:
|
|
assert in_dim == out_dim
|
|
self.layernorm = nn.LayerNorm(in_dim)
|
|
self.fc1 = nn.Linear(in_dim, hidden_dim)
|
|
self.fc2 = nn.Linear(hidden_dim, out_dim)
|
|
self.use_residual = use_residual
|
|
self.act_fn = nn.GELU()
|
|
|
|
def forward(self, x):
|
|
residual = x
|
|
x = self.layernorm(x)
|
|
x = self.fc1(x)
|
|
x = self.act_fn(x)
|
|
x = self.fc2(x)
|
|
if self.use_residual:
|
|
x = x + residual
|
|
return x |