Restoring components from modular model
- Dominant language
- Python
- Stars
- 7k
- Forks
- 2.3k
- PR merge metrics
- No merged PRs in 30d
Description
In finetuning a pretrained model, to load the pretrained *components* only, the `restore` method does not work since it fails in restoring the optimizer. It fails there because the optimizers is initialized with different parameters for the pretrainining and finetuning model.
To reproduce,
```py
from deepchem.models.torch_models import infograph
model_dir = 'model/info_star'
model = infograph.InfoGraphStarModel(3, 4, 5, model_dir=model_dir, task='supervised',
mode='regression', num_gc_layers=1)
model._ensure_built()
model.save_checkpoint()
model_ft = infograph.InfoGraphStarModel(3, 4, 5, model_dir=model_dir, task='semisupervised',
num_gc_layers=1)
model_ft._ensure_built()
model_ft.restore(components=['encoder'])
---------------------------------------------------------------------------
ValueError Traceback (most recent call last)
Cell In[26], line 16
11 model_ft = infograph.InfoGraphStarModel(3, 4, 5, model_dir=model_dir, task='semisupervised',
12 num_gc_layers=1)
14 model_ft._ensure_built()
---> 16 model_ft.restore(components=['encoder'])
File ~/proj/deepchem/models/torch_models/modular.py:401, in ModularTorchModel.restore(self, components, checkpoint, model_dir)
398 self.components[name].load_state_dict(state_dict)
400 self.build_model()
--> 401 self._pytorch_optimizer.load_state_dict(data['optimizer_state_dict'])
402 self._global_step = data['global_step']
File ~/Applications/miniconda3/envs/dc/lib/python3.9/site-packages/torch/optim/optimizer.py:390, in Optimizer.load_state_dict(self, state_dict)
388 saved_lens = (len(g['params']) for g in saved_groups)
389 if any(p_len != s_len for p_len, s_len in zip(param_lens, saved_lens)):
--> 390 raise ValueError("loaded state dict contains a parameter group "
391 "that doesn't match the size of optimizer's group")
393 # Update the state
394 id_map = {old_id: p for old_id, p in
395 zip(chain.from_iterable((g['params'] for g in saved_groups)),
396 chain.from_iterable((g['params'] for g in groups)))}
ValueError: loaded state dict contains a parameter group that doesn't match the size of optimizer's group
```
Some possible solutions:
- we can skip restoring optimizers for modular model (or)
- we can control for task and restore optimizer if the task in both instances are same (or)
- create a new method to restore components
Contributor guide
Assessment
This issue has not been assessed yet.