Skip to content

Commit

Permalink
readme
Browse files Browse the repository at this point in the history
  • Loading branch information
lucidrains authored Mar 10, 2021
1 parent 173e07e commit 7848904
Showing 1 changed file with 11 additions and 11 deletions.
22 changes: 11 additions & 11 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -150,15 +150,17 @@ import torch
from vit_pytorch import ViT
from vit_pytorch.mpp import MPP

model = ViT(image_size=256,
patch_size=32,
num_classes=1000,
dim=1024,
depth=6,
heads=8,
mlp_dim=2048,
dropout=0.1,
emb_dropout=0.1)
model = ViT(
image_size=256,
patch_size=32,
num_classes=1000,
dim=1024,
depth=6,
heads=8,
mlp_dim=2048,
dropout=0.1,
emb_dropout=0.1
)

mpp_trainer = MPP(
transformer=model,
Expand All @@ -171,11 +173,9 @@ mpp_trainer = MPP(

opt = torch.optim.Adam(mpp_trainer.parameters(), lr=3e-4)


def sample_unlabelled_images():
return torch.randn(20, 3, 256, 256)


for _ in range(100):
images = sample_unlabelled_images()
loss = mpp_trainer(images)
Expand Down

0 comments on commit 7848904

Please sign in to comment.