in projects/deep_video_compression/dvc_module.py [0:0]
def configure_optimizers(self):
# we have to train the model and the entropy bottleneck quantiles
# separately
model_param_dict, quantile_param_dict = self.model.collect_parameters()
base_optim = optim.Adam(model_param_dict.values(), lr=self.learning_rate)
if self.lr_scheduler_params is not None:
scheduler = optim.lr_scheduler.StepLR(
base_optim, **self.lr_scheduler_params
)
optimizers = [
{
"optimizer": base_optim,
"lr_scheduler": scheduler,
}
]
else:
optimizers = [{"optimizer": base_optim}]
if quantile_param_dict is not None:
optimizers.append(
{
"optimizer": optim.Adam(
quantile_param_dict.values(),
lr=self.aux_learning_rate,
)
}
)
return optimizers