Skip to content

Some questions about decoder position embedding for masked tokens #173

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Some questions about decoder position embedding for masked tokens #173

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 7 additions & 1 deletion vit_pytorch/mae.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ def __init__(
# extract some hyperparameters and functions from encoder (vision transformer to be trained)

self.encoder = encoder
# Note: This 'num_patches' contains the actual number of patches & 1 cls_token
num_patches, encoder_dim = encoder.pos_embedding.shape[-2:]
self.to_patch, self.patch_to_emb = encoder.to_patch_embedding[:2]
pixel_values_per_patch = self.patch_to_emb.weight.shape[-1]
Expand All @@ -32,6 +33,7 @@ def __init__(
self.enc_to_dec = nn.Linear(encoder_dim, decoder_dim) if encoder_dim != decoder_dim else nn.Identity()
self.mask_token = nn.Parameter(torch.randn(decoder_dim))
self.decoder = Transformer(dim = decoder_dim, depth = decoder_depth, heads = decoder_heads, dim_head = decoder_dim_head, mlp_dim = decoder_dim * 4)
# This embedding matrix also consider the ViT's cls_token
self.decoder_pos_emb = nn.Embedding(num_patches, decoder_dim)
self.to_pixels = nn.Linear(decoder_dim, pixel_values_per_patch)

Expand All @@ -41,11 +43,13 @@ def forward(self, img):
# get patches

patches = self.to_patch(img)
# Note: This 'num_patches' is the actual number of patches
batch, num_patches, *_ = patches.shape

# patch to encoder tokens and add positions

tokens = self.patch_to_emb(patches)
# pos_embedding[:, 0] is for ViT's cls_token, so we begin from 1 here
tokens = tokens + self.encoder.pos_embedding[:, 1:(num_patches + 1)]

# calculate of patches needed to be masked, and get random indices, dividing it up for mask vs unmasked
Expand Down Expand Up @@ -74,7 +78,9 @@ def forward(self, img):
# repeat mask tokens for number of masked, and add the positions using the masked indices derived above

mask_tokens = repeat(self.mask_token, 'd -> b n d', b = batch, n = num_masked)
mask_tokens = mask_tokens + self.decoder_pos_emb(masked_indices)
# Like encoder position embedding, 0 is for cls_token, so we should shift 1 here
# mask_tokens = mask_tokens + self.decoder_pos_emb(masked_indices)
mask_tokens = mask_tokens + self.decoder_pos_emb(masked_indices + 1)

# concat the masked tokens to the decoder tokens and attend with decoder

Expand Down