| import torch | |
| class CombinedOptimizer: | |
| """ | |
| A wrapper class for managing multiple optimizers. | |
| This allows different parameter groups to be optimized by different algorithms. | |
| """ | |
| def __init__(self, optimizers): | |
| self.optimizers = optimizers | |
| # Expose param_groups for learning rate schedulers | |
| self.param_groups = [] | |
| for opt in self.optimizers: | |
| self.param_groups.extend(opt.param_groups) | |
| def zero_grad(self, set_to_none=True): | |
| """Calls zero_grad() for all optimizers.""" | |
| for opt in self.optimizers: | |
| opt.zero_grad(set_to_none=set_to_none) | |
| def step(self): | |
| """Calls step() for all optimizers.""" | |
| for opt in self.optimizers: | |
| opt.step() | |
| def state_dict(self): | |
| """Returns a list of state_dicts from all optimizers.""" | |
| return [opt.state_dict() for opt in self.optimizers] | |
| def load_state_dict(self, state_dicts): | |
| """Loads state_dicts into each respective optimizer.""" | |
| for opt, state in zip(self.optimizers, state_dicts): | |
| opt.load_state_dict(state) |