src/efficientvit/apps/trainer/run_config.py

# EfficientViT: Multi-Scale Linear Attention for High-Resolution Dense Prediction
# Han Cai, Junyan Li, Muyan Hu, Chuang Gan, Song Han
# International Conference on Computer Vision (ICCV), 2023

import json

import numpy as np
import torch.nn as nn

from src.efficientvit.apps.utils import CosineLRwithWarmup, build_optimizer

__all__ = ["Scheduler", "RunConfig"]


class Scheduler:
 PROGRESS = 0


class RunConfig:
 n_epochs: int
 init_lr: float
 warmup_epochs: int
 warmup_lr: float
 lr_schedule_name: str
 lr_schedule_param: dict
 optimizer_name: str
 optimizer_params: dict
 weight_decay: float
 no_wd_keys: list
 grad_clip: float # allow none to turn off grad clipping
 reset_bn: bool
 reset_bn_size: int
 reset_bn_batch_size: int
 eval_image_size: list # allow none to use image_size in data_provider

 @property
 def none_allowed(self):
 return ["grad_clip", "eval_image_size"]

 def __init__(self, **kwargs): # arguments must be passed as kwargs
 for k, val in kwargs.items():
 setattr(self, k, val)

 # check that all relevant configs are there
 annotations = {}
 for clas in type(self).mro():
 if hasattr(clas, "__annotations__"):
 annotations.update(clas.__annotations__)
 for k, k_type in annotations.items():
 assert hasattr(
 self, k
 ), f"Key {k} with type {k_type} required for initialization."
 attr = getattr(self, k)
 if k in self.none_allowed:
 k_type = (k_type, type(None))
 assert isinstance(
 attr, k_type
 ), f"Key {k} must be type {k_type}, provided={attr}."

 self.global_step = 0
 self.batch_per_epoch = 1

 def build_optimizer(self, network: nn.Module) -> tuple[any, any]:
 r"""require setting 'batch_per_epoch' before building optimizer & lr_scheduler"""
 param_dict = {}
 for name, param in network.named_parameters():
 if param.requires_grad:
 opt_config = [self.weight_decay, self.init_lr]
 if self.no_wd_keys is not None and len(self.no_wd_keys) > 0:
 if np.any([key in name for key in self.no_wd_keys]):
 opt_config[0] = 0
 opt_key = json.dumps(opt_config)
 param_dict[opt_key] = param_dict.get(opt_key, []) + [param]

 net_params = []
 for opt_key, param_list in param_dict.items():
 wd, lr = json.loads(opt_key)
 net_params.append({"params": param_list, "weight_decay": wd, "lr": lr})

 optimizer = build_optimizer(
 net_params, self.optimizer_name, self.optimizer_params, self.init_lr
 )
 # build lr scheduler
 if self.lr_schedule_name == "cosine":
 decay_steps = []
 for epoch in self.lr_schedule_param.get("step", []):
 decay_steps.append(epoch * self.batch_per_epoch)
 decay_steps.append(self.n_epochs * self.batch_per_epoch)
 decay_steps.sort()
 lr_scheduler = CosineLRwithWarmup(
 optimizer,
 self.warmup_epochs * self.batch_per_epoch,
 self.warmup_lr,
 decay_steps,
 )
 else:
 raise NotImplementedError
 return optimizer, lr_scheduler

 def update_global_step(self, epoch, batch_id=0) -> None:
 self.global_step = epoch * self.batch_per_epoch + batch_id
 Scheduler.PROGRESS = self.progress

 @property
 def progress(self) -> float:
 warmup_steps = self.warmup_epochs * self.batch_per_epoch
 steps = max(0, self.global_step - warmup_steps)
 return steps / (self.n_epochs * self.batch_per_epoch)

 def step(self) -> None:
 self.global_step += 1
 Scheduler.PROGRESS = self.progress

 def get_remaining_epoch(self, epoch, post=True) -> int:
 return self.n_epochs + self.warmup_epochs - epoch - int(post)

 def epoch_format(self, epoch: int) -> str:
 epoch_format = f"%.{len(str(self.n_epochs))}d"
 epoch_format = f"[{epoch_format}/{epoch_format}]"
 epoch_format = epoch_format % (epoch + 1 - self.warmup_epochs, self.n_epochs)
 return epoch_format