wehub-resource-sync
133 lines
5.5 KiB
Python
133 lines
5.5 KiB
Python
import torch
|
|
from torch import Tensor
|
|
from torch.optim import Optimizer
|
|
from torch.optim.optimizer import ParamsT
|
|
from dataclasses import dataclass
|
|
from typing import Any, Dict, List, Type, Callable, Optional
|
|
|
|
|
|
@dataclass
|
|
class OptimizerSpec:
|
|
"""Spec for creating an optimizer that is part of a `ChainedOptimizer`."""
|
|
|
|
class_type: Type[Optimizer]
|
|
init_args: Dict[str, Any]
|
|
param_filter: Optional[Callable[[Tensor], bool]]
|
|
|
|
|
|
class ChainedOptimizer(Optimizer):
|
|
"""
|
|
A wrapper around multiple optimizers that allows for chaining them together.
|
|
The optimizers are applied in the order they are passed in the constructor.
|
|
Each optimizer is responsible for updating a subset of the parameters, which
|
|
is determined by the `param_filter` function. If no optimizer is found for a
|
|
parameter group, an exception is raised.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
params: ParamsT,
|
|
optimizer_specs: List[OptimizerSpec],
|
|
lr: float,
|
|
weight_decay: float = 0.0,
|
|
optimizer_selection_callback: Optional[Callable[[Tensor, int], None]] = None,
|
|
**common_kwargs,
|
|
):
|
|
self.optimizer_specs = optimizer_specs
|
|
self.optimizer_selection_callback = optimizer_selection_callback
|
|
self.optimizers: List[Optimizer] = []
|
|
defaults = dict(lr=lr, weight_decay=weight_decay)
|
|
super().__init__(params, defaults)
|
|
|
|
# Split the params for each optimizer
|
|
params_for_optimizers = [[] for _ in optimizer_specs]
|
|
for param_group in self.param_groups:
|
|
params = param_group["params"]
|
|
indices = param_group["optimizer_and_param_group_indices"] = set()
|
|
for param in params:
|
|
assert isinstance(param, Tensor), f"Expected a Tensor, got {type(param)}"
|
|
found_optimizer = False
|
|
for index, spec in enumerate(optimizer_specs):
|
|
if spec.param_filter is None or spec.param_filter(param):
|
|
if self.optimizer_selection_callback is not None:
|
|
self.optimizer_selection_callback(param, index)
|
|
params_for_optimizers[index].append(param)
|
|
indices.add((index, 0))
|
|
found_optimizer = True
|
|
break
|
|
if not found_optimizer:
|
|
raise ValueError("No valid optimizer found for the given parameter")
|
|
|
|
# Initialize the optimizers
|
|
for spec, selected_params in zip(optimizer_specs, params_for_optimizers):
|
|
optimizer_args = {
|
|
'lr': lr,
|
|
'weight_decay': weight_decay,
|
|
}
|
|
optimizer_args.update(common_kwargs)
|
|
optimizer_args.update(spec.init_args)
|
|
optimizer = spec.class_type(selected_params, **optimizer_args)
|
|
self.optimizers.append(optimizer)
|
|
|
|
def state_dict(self) -> Dict[str, Any]:
|
|
return {
|
|
"optimizers": [opt.state_dict() for opt in self.optimizers],
|
|
**super().state_dict(),
|
|
}
|
|
|
|
def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
|
|
optimizers = state_dict.pop("optimizers")
|
|
super().load_state_dict(state_dict)
|
|
for i in range(len(self.optimizers)):
|
|
self.optimizers[i].load_state_dict(optimizers[i])
|
|
|
|
def zero_grad(self, set_to_none: bool = True) -> None:
|
|
for opt in self.optimizers:
|
|
opt.zero_grad(set_to_none=set_to_none)
|
|
|
|
def _copy_lr_to_optimizers(self) -> None:
|
|
for param_group in self.param_groups:
|
|
indices = param_group["optimizer_and_param_group_indices"]
|
|
for optimizer_idx, param_group_idx in indices:
|
|
self.optimizers[optimizer_idx].param_groups[param_group_idx]["lr"] = param_group["lr"]
|
|
|
|
def step(self, closure=None) -> None:
|
|
loss = None
|
|
if closure is not None:
|
|
with torch.enable_grad():
|
|
loss = closure()
|
|
self._copy_lr_to_optimizers()
|
|
for opt in self.optimizers:
|
|
opt.step(closure=None)
|
|
return loss
|
|
|
|
def add_param_group(self, param_group: Dict[str, Any]) -> None:
|
|
super().add_param_group(param_group)
|
|
|
|
# If optimizer has not been initialized, skip adding the param groups
|
|
if not self.optimizers:
|
|
return
|
|
|
|
# Split the params for each optimizer
|
|
params_for_optimizers = [[] for _ in self.optimizer_specs]
|
|
params = param_group["params"]
|
|
indices = param_group["optimizer_and_param_group_indices"] = set()
|
|
for param in params:
|
|
assert isinstance(param, Tensor), f"Expected a Tensor, got {type(param)}"
|
|
found_optimizer = False
|
|
for index, spec in enumerate(self.optimizer_specs):
|
|
if spec.param_filter is None or spec.param_filter(param):
|
|
if self.optimizer_selection_callback is not None:
|
|
self.optimizer_selection_callback(param, index)
|
|
params_for_optimizers[index].append(param)
|
|
indices.add((index, len(self.optimizers[index].param_groups)))
|
|
found_optimizer = True
|
|
break
|
|
if not found_optimizer:
|
|
raise ValueError("No valid optimizer found for the given parameter group")
|
|
|
|
# Add the selected param group to the optimizers
|
|
for optimizer, selected_params in zip(self.optimizers, params_for_optimizers):
|
|
if selected_params:
|
|
optimizer.add_param_group({"params": selected_params})
|