class HCXVisionCAbstractor(nn.Module):
"""
This module is based on C-Abstractor, whose license is under apache-2.0.
You can check the original code at
https://github.com/khanrc/honeybee/blob/main/honeybee/projectors/projectors.py
and we made necessary modifications.
"""
def __init__(
self,
num_queries: int,
num_input_tokens: int,
encoder_hidden_size: int,
hidden_size: int,
output_hidden_size: int,
pos_emb: bool = True,
prenorm: bool = False,
):
super().__init__()
self.num_input_tokens = num_input_tokens
self.output_hidden_size = output_hidden_size
# Positional embedding
if pos_emb:
self.pos_emb = torch.nn.Parameter(
torch.zeros(1, num_input_tokens, encoder_hidden_size)
)
self.pos_emb.data.normal_(mean=0.0, std=0.02)
else:
self.pos_emb = None
# (Optional) Pre-normalization layer
if prenorm:
self.prenorm = LayerNorm(encoder_hidden_size)
else:
self.prenorm = None
self.build_net(
num_queries, encoder_hidden_size, hidden_size, output_hidden_size
)
self.dtype = next(self.parameters()).dtype
def forward(
self,
x: torch.Tensor,
num_queries_vis_abstractors: list[list[int]] | None = None,
num_grids: list[int] | None = None,
) -> torch.Tensor:
if self.prenorm is not None:
x = self.prenorm(x)
if self.pos_emb is not None:
x = x + self.pos_emb
x = self._forward(
x,
num_queries_vis_abstractors=num_queries_vis_abstractors,
num_grids=num_grids,
) # (B, L, output_hidden_size)
return x
def _forward(
self,
x: torch.Tensor,
num_queries_vis_abstractors: list[list[int]] | None = None,
num_grids: list[int] | None = None,
) -> torch.Tensor:
# x: [B, L, dim]
B, L, dim = x.shape
hw = int(L**0.5)
x = rearrange(x, "b (h w) d -> b d h w", h=hw, w=hw)
if num_queries_vis_abstractors is not None:
assert num_grids is not None
return self._forward_adaptive_num_query(
x, num_queries_vis_abstractors, num_grids
)
x = self.net(x)
x = rearrange(x, "b d h w -> b (h w) d")
x = self.readout(x)
return x
def _forward_adaptive_num_query(
self,
x: torch.Tensor,
num_queries_vis_abstractors: list[list[int]] | None = None,
num_grids: list[int] | None = None,
) -> list[torch.Tensor]:
# self.net is consisted by 3 layers (s1, sampler, s2)
assert len(self.net) == 3
x = self.net[0](x) # s1
new_x = []
for i, num_queries in enumerate(num_queries_vis_abstractors):
hw = int(num_queries**0.5)
sampler = nn.AdaptiveAvgPool2d((hw, hw))
out = sampler(x[num_grids[i] : num_grids[i + 1], :])
out = self.net[2](out) # s2
out = rearrange(out, "b d h w -> b (h w) d")
out = self.readout(out)
new_x.append(out)
return new_x
def build_net(
self,
n_queries: int,
encoder_hidden_size: int,
hidden_size: int,
output_hidden_size: int,
depth: int = 3,
mlp_depth: int = 2,
):
assert (n_queries**0.5).is_integer(), (
f"n_queries must be square number. n_queries: {n_queries}"
)
hw = int(n_queries**0.5)
# RegBlock = ResBlock + SE
RegBlock = partial(
RegStage,
stride=1,
dilation=1,
act_layer=nn.SiLU,
norm_layer=LayerNorm2d,
)
s1 = RegBlock(
depth,
encoder_hidden_size,
hidden_size,
)
sampler = nn.AdaptiveAvgPool2d((hw, hw))
s2 = RegBlock(
depth,
hidden_size,
hidden_size,
)
self.net = nn.Sequential(s1, sampler, s2)
self.readout = self.build_mlp(mlp_depth, hidden_size, output_hidden_size)
def build_mlp(
self,
depth: int,
hidden_size: int,
output_hidden_size: int,
):
layers = [nn.Linear(hidden_size, output_hidden_size)]
for _ in range(1, depth):
layers.append(nn.SiLU())
layers.append(nn.Linear(output_hidden_size, output_hidden_size))
return nn.Sequential(*layers)