add lars opt and configs

src/train_supervised.py

@@ -16,6 +16,7 @@ from src.models.head import ViTClassifier
 from src.transforms import make_transforms, make_transform_eval
 from src.utils.distributed import init_distributed
 from src.utils.logging import CSVLogger, AverageMeter
+from src.utils.optimizers import LARS
 
 # --
 log_freq = 10
@@ -217,17 +218,38 @@ def main(args, resume_preempt=False):
         logger.info("Training full model (Fine-tuning mode)")
 
     params = [p for p in model.parameters() if p.requires_grad]
-    if o_args["optimizer"].lower() == "adamw":
+    optimizer_name = o_args["optimizer"].lower()
+    if optimizer_name == "adamw":
         optimizer = torch.optim.AdamW(
             params, lr=o_args["lr"], weight_decay=o_args["weight_decay"]
         )
+    elif optimizer_name == "lars":
+        optimizer = LARS(
+            params,
+            lr=o_args["lr"],
+            weight_decay=o_args["weight_decay"],
+            momentum=o_args.get("momentum", 0.9),
+            eta=o_args.get("lars_eta", 0.001),
+            eps=o_args.get("lars_eps", 1e-8),
+            exclude_bias_and_norm=o_args.get("lars_exclude_bias_and_norm", True),
+        )
     else:
         optimizer = torch.optim.SGD(
-            params, lr=o_args["lr"], momentum=0.9, weight_decay=o_args["weight_decay"]
+            params,
+            lr=o_args["lr"],
+            momentum=o_args.get("momentum", 0.9),
+            weight_decay=o_args["weight_decay"],
         )
 
     scheduler = None
-    if o_args["use_cosine_schedule"]:
+    lr_schedule = o_args.get("lr_schedule", "cosine").lower()
+    if lr_schedule == "step":
+        scheduler = torch.optim.lr_scheduler.MultiStepLR(
+            optimizer,
+            milestones=o_args.get("step_milestones", [15, 30, 45]),
+            gamma=o_args.get("step_gamma", 0.1),
+        )
+    elif o_args["use_cosine_schedule"]:
         scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
             optimizer, T_max=o_args["epochs"], eta_min=o_args["final_lr"]
         )

src/utils/optimizers.py

@@ -0,0 +1,84 @@
+import math
+
+import torch
+
+
+class LARS(torch.optim.Optimizer):
+    def __init__(
+        self,
+        params,
+        lr,
+        weight_decay=0.0,
+        momentum=0.9,
+        eta=0.001,
+        eps=1e-8,
+        exclude_bias_and_norm=True,
+    ):
+        if lr <= 0.0:
+            raise ValueError(f"Invalid lr: {lr}")
+        if weight_decay < 0.0:
+            raise ValueError(f"Invalid weight_decay: {weight_decay}")
+        if momentum < 0.0:
+            raise ValueError(f"Invalid momentum: {momentum}")
+        if eta <= 0.0:
+            raise ValueError(f"Invalid eta: {eta}")
+        if eps <= 0.0:
+            raise ValueError(f"Invalid eps: {eps}")
+
+        defaults = dict(
+            lr=lr,
+            weight_decay=weight_decay,
+            momentum=momentum,
+            eta=eta,
+            eps=eps,
+            exclude_bias_and_norm=exclude_bias_and_norm,
+        )
+        super().__init__(params, defaults)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            lr = group["lr"]
+            weight_decay = group["weight_decay"]
+            momentum = group["momentum"]
+            eta = group["eta"]
+            eps = group["eps"]
+            exclude_bias_and_norm = group["exclude_bias_and_norm"]
+
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+
+                grad = p.grad
+                if grad.is_sparse:
+                    raise RuntimeError("LARS does not support sparse gradients")
+
+                param_norm = torch.norm(p)
+                grad_norm = torch.norm(grad)
+
+                lars_lr = 1.0
+                if not exclude_bias_and_norm or p.ndim > 1:
+                    if param_norm > 0.0 and grad_norm > 0.0:
+                        lars_lr = eta * param_norm / (grad_norm + weight_decay * param_norm + eps)
+
+                d_p = grad
+                if weight_decay != 0.0 and (not exclude_bias_and_norm or p.ndim > 1):
+                    d_p = d_p.add(p, alpha=weight_decay)
+
+                if momentum != 0.0:
+                    param_state = self.state.setdefault(p, {})
+                    if "momentum_buffer" not in param_state:
+                        buf = param_state["momentum_buffer"] = torch.clone(d_p).detach()
+                    else:
+                        buf = param_state["momentum_buffer"]
+                        buf.mul_(momentum).add_(d_p)
+                    d_p = buf
+
+                p.add_(d_p, alpha=-lr * lars_lr)
+
+        return loss