File size: 3,420 Bytes
eedeabf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 |
import torch
from einops import rearrange
from torch import nn, Tensor
from torch.nn import LayerNorm, Linear, ModuleList
from .modules import Block, no_grad_trunc_normal_
from .positional_embedding import SinCosPositionalEmbedding
class MarlinDecoder(nn.Module):
def __init__(self, img_size=224, patch_size=16, n_frames=16, embed_dim=384, depth=8,
num_heads=6, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0.,
norm_layer="LayerNorm", init_values=1., tubelet_size=2
):
super().__init__()
output_dim = 3 * tubelet_size * patch_size * patch_size
self.patch_size = patch_size
self.tubelet_size = tubelet_size
self.n_patch_h = img_size // patch_size
self.n_patch_w = img_size // patch_size
self.embed_dim = embed_dim
if norm_layer == "LayerNorm":
self.norm_layer = LayerNorm
self.norm = self.norm_layer(embed_dim)
else:
raise NotImplementedError("Only LayerNorm is supported")
# sine-cosine positional embeddings
self.pos_embedding = SinCosPositionalEmbedding(
(self.n_patch_h * self.n_patch_w * (n_frames // tubelet_size), embed_dim), dropout_rate=0.)
self.mask_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
self.blocks = ModuleList([
Block(
dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
drop=drop_rate, attn_drop=attn_drop_rate, norm_layer=self.norm_layer,
init_values=init_values
) for _ in range(depth)])
self.head = Linear(embed_dim, output_dim)
self.apply(self._init_weights)
no_grad_trunc_normal_(self.mask_token, mean=0., std=0.02, a=-0.02, b=0.02)
@staticmethod
def _init_weights(m):
if isinstance(m, nn.Linear):
nn.init.xavier_uniform_(m.weight)
if isinstance(m, nn.Linear) and m.bias is not None:
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.LayerNorm):
nn.init.constant_(m.bias, 0)
nn.init.constant_(m.weight, 1.0)
def unpatch_to_img(self, x: Tensor) -> Tensor:
# x: (Batch, No. batches, Prod of cube size * C)
x = rearrange(x, "b n (c p) -> b n p c", c=3)
# x: (Batch, No. batches, Prod of cube size, C)
x = rearrange(x, "b (t h w) (p0 p1 p2) c -> b c (t p0) (h p1) (w p2)", p0=self.tubelet_size,
p1=self.patch_size, p2=self.patch_size, h=self.n_patch_h, w=self.n_patch_w)
# x: (B, C, T, H, W)
return x
def forward_features(self, x, return_token_num=0):
for block in self.blocks:
x = block(x)
if return_token_num > 0:
x = x[:, -return_token_num:]
x = self.norm(x)
x = self.head(x)
# x: (B, N_mask, C)
return x
def forward(self, x, mask):
# mask: 0 -> masked, 1 -> visible
b, n, c = x.shape
expand_pos_embed = self.pos_embedding.emb.data.expand(b, -1, -1)
pos_emb_vis = expand_pos_embed[mask].view(b, -1, c)
pos_emb_mask = expand_pos_embed[~mask].view(b, -1, c)
x = torch.cat([x + pos_emb_vis, self.mask_token + pos_emb_mask], dim=1)
mask_num = pos_emb_mask.shape[1]
x = self.forward_features(x, return_token_num=mask_num)
return x
|