deepchem / deepchem/deepchem

Restoring components from modular model

Open
#3,490 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.