learning/annotated_deep_learning_paper_implementations
1
0
Fork 0
mirror of https://github.com/labmlai/annotated_deep_learning_paper_implementations.git synced 2025-08-17 11:20:32 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: fix cls_token bug in vit.

Browse Source
This commit is contained in:
Li zhuoqun
2023-11-06 11:28:45 +08:00
parent 9a42ac2697
commit ffafaf1df7
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
labml_nn/transformers/vit/__init__.py
View File

@ -191,11 +191,11 @@ class VisionTransformer(Module):
""" """
# Get patch embeddings. This gives a tensor of shape `[patches, batch_size, d_model]` # Get patch embeddings. This gives a tensor of shape `[patches, batch_size, d_model]`
x = self.patch_emb(x) x = self.patch_emb(x)
# Add positional embeddings
x = self.pos_emb(x)
# Concatenate the `[CLS]` token embeddings before feeding the transformer # Concatenate the `[CLS]` token embeddings before feeding the transformer
cls_token_emb = self.cls_token_emb.expand(-1, x.shape[1], -1) cls_token_emb = self.cls_token_emb.expand(-1, x.shape[1], -1)
x = torch.cat([cls_token_emb, x]) x = torch.cat([cls_token_emb, x])
# Add positional embeddings
x = self.pos_emb(x)
# Pass through transformer layers with no attention masking # Pass through transformer layers with no attention masking
for layer in self.transformer_layers: for layer in self.transformer_layers: