Spaces:
Running
Running
import math | |
from typing import Tuple, Type | |
import numpy as np | |
import torch | |
from torch import nn, Tensor | |
# Lightly adapted from | |
# https://github.com/facebookresearch/MaskFormer/blob/main/mask_former/modeling/transformer/transformer_predictor.py # noqa | |
class MLPBlock(nn.Module): | |
def __init__( | |
self, | |
input_dim: int, | |
hidden_dim: int, | |
output_dim: int, | |
num_layers: int, | |
act: Type[nn.Module], | |
) -> None: | |
super().__init__() | |
self.num_layers = num_layers | |
h = [hidden_dim] * (num_layers - 1) | |
self.layers = nn.ModuleList( | |
nn.Sequential(nn.Linear(n, k), act()) | |
for n, k in zip([input_dim] + h, [hidden_dim] * num_layers) | |
) | |
self.fc = nn.Linear(hidden_dim, output_dim) | |
def forward(self, x): | |
for layer in self.layers: | |
x = layer(x) | |
return self.fc(x) | |
# From https://github.com/yformer/EfficientSAM/blob/main/efficient_sam/efficient_sam_decoder.py | |
class PositionEmbeddingRandom(nn.Module): | |
""" | |
Positional encoding using random spatial frequencies. | |
""" | |
def __init__(self, num_pos_feats: int) -> None: | |
super().__init__() | |
self.register_buffer( | |
"positional_encoding_gaussian_matrix", torch.randn((2, num_pos_feats)) | |
) | |
def _pe_encoding(self, coords: torch.Tensor) -> torch.Tensor: | |
"""Positionally encode points that are normalized to [0,1].""" | |
# assuming coords are in [0, 1]^2 square and have d_1 x ... x d_n x 2 shape | |
coords = 2 * coords - 1 | |
coords = coords @ self.positional_encoding_gaussian_matrix | |
coords = 2 * np.pi * coords | |
# outputs d_1 x ... x d_n x C shape | |
return torch.cat([torch.sin(coords), torch.cos(coords)], dim=-1) | |
def forward(self, size: Tuple[int, int]) -> torch.Tensor: | |
"""Generate positional encoding for a grid of the specified size.""" | |
h, w = size | |
device = self.positional_encoding_gaussian_matrix.device | |
grid = torch.ones([h, w], device=device, dtype=torch.float32) | |
y_embed = grid.cumsum(dim=0) - 0.5 | |
x_embed = grid.cumsum(dim=1) - 0.5 | |
y_embed = y_embed / h | |
x_embed = x_embed / w | |
pe = self._pe_encoding(torch.stack([x_embed, y_embed], dim=-1)) | |
return pe.permute(2, 0, 1) # C x H x W | |
def forward_with_coords( | |
self, coords_input: torch.Tensor, image_size: Tuple[int, int] | |
) -> torch.Tensor: | |
"""Positionally encode points that are not normalized to [0,1].""" | |
coords = coords_input.clone() | |
coords[:, :, 0] = coords[:, :, 0] / image_size[1] | |
coords[:, :, 1] = coords[:, :, 1] / image_size[0] | |
return self._pe_encoding(coords.to(torch.float)) # B x N x C | |
# From https://github.com/yformer/EfficientSAM/blob/main/efficient_sam/build_efficient_sam.py | |
class TwoWayTransformer(nn.Module): | |
def __init__( | |
self, | |
depth: int, | |
embedding_dim: int, | |
num_heads: int, | |
mlp_dim: int, | |
activation: Type[nn.Module], | |
normalize_before_activation: bool, | |
attention_downsample_rate: int = 2, | |
) -> None: | |
""" | |
A transformer decoder that attends to an input image using | |
queries whose positional embedding is supplied. | |
Args: | |
depth (int): number of layers in the transformer | |
embedding_dim (int): the channel dimension for the input embeddings | |
num_heads (int): the number of heads for multihead attention. Must | |
divide embedding_dim | |
mlp_dim (int): the channel dimension internal to the MLP block | |
activation (nn.Module): the activation to use in the MLP block | |
""" | |
super().__init__() | |
self.depth = depth | |
self.embedding_dim = embedding_dim | |
self.num_heads = num_heads | |
self.mlp_dim = mlp_dim | |
self.layers = nn.ModuleList() | |
for i in range(depth): | |
curr_layer = TwoWayAttentionBlock( | |
embedding_dim=embedding_dim, | |
num_heads=num_heads, | |
mlp_dim=mlp_dim, | |
activation=activation, | |
normalize_before_activation=normalize_before_activation, | |
attention_downsample_rate=attention_downsample_rate, | |
skip_first_layer_pe=(i == 0), | |
) | |
self.layers.append(curr_layer) | |
self.final_attn_token_to_image = AttentionForTwoWayAttentionBlock( | |
embedding_dim, | |
num_heads, | |
downsample_rate=attention_downsample_rate, | |
) | |
self.norm_final_attn = nn.LayerNorm(embedding_dim) | |
def forward( | |
self, | |
image_embedding: Tensor, | |
image_pe: Tensor, | |
point_embedding: Tensor, | |
) -> Tuple[Tensor, Tensor]: | |
""" | |
Args: | |
image_embedding (torch.Tensor): image to attend to. Should be shape | |
B x embedding_dim x h x w for any h and w. | |
image_pe (torch.Tensor): the positional encoding to add to the image. Must | |
have the same shape as image_embedding. | |
point_embedding (torch.Tensor): the embedding to add to the query points. | |
Must have shape B x N_points x embedding_dim for any N_points. | |
Returns: | |
torch.Tensor: the processed point_embedding | |
torch.Tensor: the processed image_embedding | |
""" | |
# BxCxHxW -> BxHWxC == B x N_image_tokens x C | |
# bs, c, h, w = image_embedding.shape | |
if len(image_embedding.shape) == 4: | |
image_embedding = image_embedding.flatten(2).permute(0, 2, 1) | |
image_pe = image_pe.flatten(2).permute(0, 2, 1) | |
# Prepare queries | |
queries = point_embedding | |
keys = image_embedding | |
# Apply transformer blocks and final layernorm | |
for idx, layer in enumerate(self.layers): | |
queries, keys = layer( | |
queries=queries, | |
keys=keys, | |
query_pe=point_embedding, | |
key_pe=image_pe, | |
) | |
# Apply the final attention layer from the points to the image | |
q = queries + point_embedding | |
k = keys + image_pe | |
attn_out = self.final_attn_token_to_image(q=q, k=k, v=keys) | |
queries = queries + attn_out | |
queries = self.norm_final_attn(queries) | |
return queries, keys | |
class TwoWayAttentionBlock(nn.Module): | |
def __init__( | |
self, | |
embedding_dim: int, | |
num_heads: int, | |
mlp_dim: int, | |
activation: Type[nn.Module], | |
normalize_before_activation: bool, | |
attention_downsample_rate: int = 2, | |
skip_first_layer_pe: bool = False, | |
) -> None: | |
""" | |
A transformer block with four layers: (1) self-attention of sparse | |
inputs, (2) cross attention of sparse inputs to dense inputs, (3) mlp | |
block on sparse inputs, and (4) cross attention of dense inputs to sparse | |
inputs. | |
Arguments: | |
embedding_dim (int): the channel dimension of the embeddings | |
num_heads (int): the number of heads in the attention layers | |
mlp_dim (int): the hidden dimension of the mlp block | |
activation (nn.Module): the activation of the mlp block | |
skip_first_layer_pe (bool): skip the PE on the first layer | |
""" | |
super().__init__() | |
self.self_attn = AttentionForTwoWayAttentionBlock(embedding_dim, num_heads) | |
self.norm1 = nn.LayerNorm(embedding_dim) | |
self.cross_attn_token_to_image = AttentionForTwoWayAttentionBlock( | |
embedding_dim, | |
num_heads, | |
downsample_rate=attention_downsample_rate, | |
) | |
self.norm2 = nn.LayerNorm(embedding_dim) | |
self.mlp = MLPBlock( | |
embedding_dim, | |
mlp_dim, | |
embedding_dim, | |
1, | |
activation, | |
) | |
self.norm3 = nn.LayerNorm(embedding_dim) | |
self.norm4 = nn.LayerNorm(embedding_dim) | |
self.cross_attn_image_to_token = AttentionForTwoWayAttentionBlock( | |
embedding_dim, | |
num_heads, | |
downsample_rate=attention_downsample_rate, | |
) | |
self.skip_first_layer_pe = skip_first_layer_pe | |
def forward( | |
self, queries: Tensor, keys: Tensor, query_pe: Tensor, key_pe: Tensor | |
) -> Tuple[Tensor, Tensor]: | |
# Self attention block | |
if not self.skip_first_layer_pe: | |
queries = queries + query_pe | |
attn_out = self.self_attn(q=queries, k=queries, v=queries) | |
queries = queries + attn_out | |
queries = self.norm1(queries) | |
# Cross attention block, tokens attending to image embedding | |
q = queries + query_pe | |
k = keys + key_pe | |
attn_out = self.cross_attn_token_to_image(q=q, k=k, v=keys) | |
queries = queries + attn_out | |
queries = self.norm2(queries) | |
# MLP block | |
mlp_out = self.mlp(queries) | |
queries = queries + mlp_out | |
queries = self.norm3(queries) | |
# Cross attention block, image embedding attending to tokens | |
q = queries + query_pe | |
k = keys + key_pe | |
attn_out = self.cross_attn_image_to_token(q=k, k=q, v=queries) | |
keys = keys + attn_out | |
keys = self.norm4(keys) | |
return queries, keys | |
class AttentionForTwoWayAttentionBlock(nn.Module): | |
""" | |
An attention layer that allows for downscaling the size of the embedding | |
after projection to queries, keys, and values. | |
""" | |
def __init__( | |
self, | |
embedding_dim: int, | |
num_heads: int, | |
downsample_rate: int = 1, | |
) -> None: | |
super().__init__() | |
self.embedding_dim = embedding_dim | |
self.internal_dim = embedding_dim // downsample_rate | |
self.num_heads = num_heads | |
assert ( | |
self.internal_dim % num_heads == 0 | |
), "num_heads must divide embedding_dim." | |
self.q_proj = nn.Linear(embedding_dim, self.internal_dim) | |
self.k_proj = nn.Linear(embedding_dim, self.internal_dim) | |
self.v_proj = nn.Linear(embedding_dim, self.internal_dim) | |
self.out_proj = nn.Linear(self.internal_dim, embedding_dim) | |
self._reset_parameters() | |
def _reset_parameters(self) -> None: | |
# The fan_out is incorrect, but matches pytorch's initialization | |
# for which qkv is a single 3*embedding_dim x embedding_dim matrix | |
fan_in = self.embedding_dim | |
fan_out = 3 * self.internal_dim | |
# Xavier uniform with our custom fan_out | |
bnd = math.sqrt(6 / (fan_in + fan_out)) | |
nn.init.uniform_(self.q_proj.weight, -bnd, bnd) | |
nn.init.uniform_(self.k_proj.weight, -bnd, bnd) | |
nn.init.uniform_(self.v_proj.weight, -bnd, bnd) | |
# out_proj.weight is left with default initialization, like pytorch attention | |
nn.init.zeros_(self.q_proj.bias) | |
nn.init.zeros_(self.k_proj.bias) | |
nn.init.zeros_(self.v_proj.bias) | |
nn.init.zeros_(self.out_proj.bias) | |
def _separate_heads(self, x: Tensor, num_heads: int) -> Tensor: | |
b, n, c = x.shape | |
x = x.reshape(b, n, num_heads, c // num_heads) | |
return x.transpose(1, 2) # B x N_heads x N_tokens x C_per_head | |
def _recombine_heads(self, x: Tensor) -> Tensor: | |
b, n_heads, n_tokens, c_per_head = x.shape | |
x = x.transpose(1, 2) | |
return x.reshape(b, n_tokens, n_heads * c_per_head) # B x N_tokens x C | |
def forward(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor: | |
# Input projections | |
q = self.q_proj(q) | |
k = self.k_proj(k) | |
v = self.v_proj(v) | |
# Separate into heads | |
q = self._separate_heads(q, self.num_heads) | |
k = self._separate_heads(k, self.num_heads) | |
v = self._separate_heads(v, self.num_heads) | |
# Attention | |
_, _, _, c_per_head = q.shape | |
attn = q @ k.permute(0, 1, 3, 2) # B x N_heads x N_tokens x N_tokens | |
attn = attn / math.sqrt(c_per_head) | |
attn = torch.softmax(attn, dim=-1) | |
# Get output | |
out = attn @ v | |
out = self._recombine_heads(out) | |
out = self.out_proj(out) | |
return out | |