def prepare_optimizer(self):
"""
Constructs the optimizer using the user defined settings in the yaml config.
The model must be on the correct device (cuda or cpu) by this point.
"""
param_groups = get_optimizer_param_groups(
model=self.base_model,
model_config=self.config["MODEL"],
optimizer_config=self.config["OPTIMIZER"],
optimizer_schedulers=self.optimizer_schedulers,
)
self.optimizer.set_param_groups(param_groups)
def _compute_param_groups(finetune_config: AttrDict):
"""
Take a configuration and compute the parameter groups
for this configuration
"""
optimizer_schedulers = build_optimizer_schedulers(
finetune_config["OPTIMIZER"])
base_model = build_model(finetune_config["MODEL"],
finetune_config["OPTIMIZER"])
return get_optimizer_param_groups(
model=base_model,
model_config=finetune_config["MODEL"],
optimizer_config=finetune_config["OPTIMIZER"],
optimizer_schedulers=optimizer_schedulers,
)