Skip to content

Commit

Permalink
decoder positional embedding needs to be reapplied https://twitter.co…
Browse files Browse the repository at this point in the history
  • Loading branch information
lucidrains committed Jan 6, 2022
1 parent 28eaba6 commit 1cc0f18
Show file tree
Hide file tree
Showing 2 changed files with 4 additions and 7 deletions.
2 changes: 1 addition & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
setup(
name = 'vit-pytorch',
packages = find_packages(exclude=['examples']),
version = '0.26.2',
version = '0.26.3',
license='MIT',
description = 'Vision Transformer (ViT) - Pytorch',
author = 'Phil Wang',
Expand Down
9 changes: 3 additions & 6 deletions vit_pytorch/mae.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,13 +14,11 @@ def __init__(
masking_ratio = 0.75,
decoder_depth = 1,
decoder_heads = 8,
decoder_dim_head = 64,
apply_decoder_pos_emb_all = False # whether to (re)apply decoder positional embedding to encoder unmasked tokens
decoder_dim_head = 64
):
super().__init__()
assert masking_ratio > 0 and masking_ratio < 1, 'masking ratio must be kept between 0 and 1'
self.masking_ratio = masking_ratio
self.apply_decoder_pos_emb_all = apply_decoder_pos_emb_all

# extract some hyperparameters and functions from encoder (vision transformer to be trained)

Expand Down Expand Up @@ -73,10 +71,9 @@ def forward(self, img):

decoder_tokens = self.enc_to_dec(encoded_tokens)

# reapply decoder position embedding to unmasked tokens, if desired
# reapply decoder position embedding to unmasked tokens

if self.apply_decoder_pos_emb_all:
decoder_tokens = decoder_tokens + self.decoder_pos_emb(unmasked_indices)
decoder_tokens = decoder_tokens + self.decoder_pos_emb(unmasked_indices)

# repeat mask tokens for number of masked, and add the positions using the masked indices derived above

Expand Down

0 comments on commit 1cc0f18

Please sign in to comment.