ZF/vit-pytorch
1
0
Fork 0
mirror of https://github.com/lucidrains/vit-pytorch.git synced 2025-12-30 08:02:29 +00:00
Code Issues Packages Projects Releases Wiki Activity

readme

Browse Source
This commit is contained in:
Phil Wang
2021-11-12 20:19:38 -08:00
parent e8f6d72033
commit 5b16e8f809
1 changed files with 6 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
README.md
View File

@@ -514,10 +514,9 @@ v = ViT(
mae = MAE( mae = MAE(
encoder = v, encoder = v,
masking_ratio = 0.75, masking_ratio = 0.75, # the paper recommended 75% masked patches
decoder_dim = 1024, decoder_dim = 512, # paper showed good results with just 512
decoder_depth = 6, decoder_depth = 6 # anywhere from 1 to 8
decoder_heads = 8
) )
images = torch.randn(8, 3, 256, 256) images = torch.randn(8, 3, 256, 256)
@@ -527,6 +526,9 @@ loss.backward()
# that's all! # that's all!
# do the above in a for loop many times with a lot of images and your vision transformer will learn # do the above in a for loop many times with a lot of images and your vision transformer will learn
# save your improved vision transformer
torch.save(v.state_dict(), './trained-vit.pt')
``` ```
## Masked Patch Prediction ## Masked Patch Prediction