Loading pretrained torch models
- 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
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