Pavel/vision (#16)

fix lr scheduler
openai · Dec 18, 2023 · 0a6d4ab · 0a6d4ab
1 parent 6ce5200
commit 0a6d4ab
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/vision/README.md b/vision/README.md
@@ -48,9 +48,9 @@ With the commands above we get the following results (note that the results may
 | Model                   | Top-1 Accuracy |
 |-------------------------|----------------|
 | AlexNet                 | 56.6           |
-| Dino ResNet50           | 63.7           |
-| Dino ViT-B/8            | 74.9           |
-| AlexNet → DINO ResNet50 | 60.7           |
-| AlexNet → DINO ViT-B/8  | 64.2           |
+| Dino ResNet50           | 64.5           |
+| Dino ViT-B/8            | 74.0           |
+| AlexNet → DINO ResNet50 | 61.9           |
+| AlexNet → DINO ViT-B/8  | 66.6           |
 
 You can add new custom models to the `models.py` and new datasets to `data.py`.
diff --git a/vision/run_weak_strong.py b/vision/run_weak_strong.py
@@ -66,7 +66,9 @@ def train_logreg(
     model = torch.nn.Linear(d, n_classes).cuda()
     criterion = torch.nn.CrossEntropyLoss()
     optimizer = torch.optim.Adam(model.parameters(), weight_decay=weight_decay, lr=lr)
-    schedule = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer=optimizer, T_max=n_epochs)
+    n_batches = len(train_loader)
+    n_iter = n_batches * n_epochs
+    schedule = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer=optimizer, T_max=n_iter)
 
     results = {f"{key}_all": [] for key in eval_datasets.keys()}
     for epoch in (pbar := tqdm.tqdm(range(n_epochs), desc="Epoch 0")):