torchwrench.optim package¶
-
class torchwrench.optim.CosDecayScheduler(optimizer: Optimizer, n_steps: int, last_epoch: int =
-1)[source]¶ Bases:
LambdaLR
-
torchwrench.optim.create_params_groups_bias(model_or_params: Module | Iterable[tuple[str, Parameter]], weight_decay: float, *, skip_list: Iterable[str] | None =
(), verbose: int =2) list[dict[str, list[Parameter] | float]][source]¶ Split parameters into 2 groups with or without weight decay for AdamW optimizer.
Example¶
` >>> model = nn.Linear(100, 10) >>> weight_decay = 0.01 >>> param_groups = create_params_groups_bias(model, weight_decay=weight_decay) >>> optimizer = AdamW(params_groups, weight_decay=weight_decay) `
-
torchwrench.optim.get_lr(optim: Optimizer, idx: int =
0, key: str ='lr') float[source]¶ Get the learning rate of the first group of an optimizer.
- Args:
optim: The optimizer to get. idx: The group index of the learning rate in the optimizer. defaults to 0.
-
torchwrench.optim.get_lrs(optim: Optimizer, key: str =
'lr') list[float][source]¶ Get the learning rates in all groups of an optimizer.
- Args:
optim: The optimizer to get.