deepchem / deepchem/deepchem

Loading pretrained torch models

Open
#3,155 0 comments 0 reactions 0 assignees View on GitHub
enhancement
Dominant language
Python
Stars
7k
Forks
2.3k
PR merge metrics
No merged PRs in 30d

Description

To load a pretrained model with only specific layers, user has to create a assignment map between source model tensor and destination model tensors. This process is bit cumbersome. For example, this is the code which I have to use in the current setup for loading a pretrained model:

```
class Embedding(nn.Module):

def __init__(self):
super(Embedding, self).__init__()
self.dense_1 = nn.Linear(in_features=16, out_features=4, bias=True)

def forward(self, x):
return self.dense_2(self.dense_1(x))

class Model(nn.Module):

def __init__(self):
super(Model, self).__init__()
self.embedding = Embedding()
self.head = nn.Linear(in_features=4, out_features=1, bias=True)

def forward(self, x):
return self.head(self.embedding(x))

pre = Model()
pre_dc = dc.models.TorchModel(pre, loss=dc.models.losses.L2Loss())
fine = Model()
fine_dc = dc.models.TorchModel(fine, loss=dc.models.losses.L2Loss())

# creating value map between source tensor and destination tensor
value_map = {}
svs = list(pre.parameters())
for sv in svs:
value_map[sv] = sv.detach().cpu().numpy()

assignment_map = {}
fvs = list(fine.parameters())

assignment_map[svs[0]] = fvs[0] # I only want to copy the embedding layer, hence I do a manual assignment based on the position of the embedding layer in the tensor.
assignment_map[svs[1]] = fvs[1] # Bias of embedding layer

fine_dc.load_from_pretrained(pre_dc, assignment_map=assignment_map, value_map=value_map)
```

Instead, if we use `state_dict` keys for assignment map, it will be more user-friendly. In this scenario, assignment map will be
```
assignment_map = {'embedding.dense_1.weight': 'embedding.dense_1.weight'}
```

and `TorchModel.load_from_pretrained` can be modified as:
```
def load_from_pretrained(source_model, dest_model, assignment_map):
new_state_dict = {}
for key in assignment_map.keys():
new_state_dict[key] = pre.state_dict()[key]
return fine.load_state_dict(new_state_dict, strict=False)
```

Contributor guide

Open the contributing guide

Research direction

Start at the TorchModel.load_from_pretrained entry point and inspect how assignment_map and value_map currently select tensors. Support state_dict key mappings for partial pretrained loading, and confirm that the destination model accepts the selected weights without requiring positional parameter mapping.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Feature
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Clearly specified
Newbie friendliness
50/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.