File size: 1,118 Bytes
ad93bc0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 | 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) |