Merge branch 'master' of https://github.com/Andrewzh112/codebase

Andrewzh112 · Dec 2, 2020 · 0ac2d0d · 0ac2d0d
2 parents 072e193 + 4cb3459
commit 0ac2d0d
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/simsiam/main.py b/simsiam/main.py
@@ -87,7 +87,7 @@ def cosine_loss(p, z):
 
     optimizer = torch.optim.SGD(model.parameters(), lr=args.lr,
                                 momentum=args.momentum, weight_decay=args.wd)
-    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, args.epochs // 40)
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, int(args.epochs * 0.005))
 
     pbar = tqdm(range(args.epochs))
     for epoch in pbar:
@@ -137,7 +137,8 @@ def cosine_loss(p, z):
         tqdm.write(
             f'Epoch {epoch + 1}/{args.epochs}, \
                 Train Loss: {sum(train_losses) / len(train_losses):.3f}, \
-                Top Acc @ 1: {top1acc:.3f}'
+                Top Acc @ 1: {top1acc:.3f}, \
+                Learning Rate: {scheduler.get_last_lr()}'
         )
         torch.save(model.state_dict(), args.check_point)
         scheduler.step()