Dear experts,
Recently I have a need to load a torch scripted model and modify the last layer to finetune.
An example is provided as below:
from torchvision import models
import torch
import torch.nn as nn
# Example of how I saved the model.
model_ft = models.resnet18(pretrained=True)
model_ft_scripted = torch.jit.script(model_ft)
model_ft_scripted.save('r18_model_ft_scripted.pt')
# Load the model and adjust the last layer.
model_ft_scripted_loaded = torch.jit.load("r18_model_ft_scripted.pt")
model_ft_scripted_loaded.fc = nn.Linear(512, 2)
And it complains Cannot re-assign modules in a ScriptModule with non-scripted module.
Then I tried:
linear_module = nn.Linear(512, 2)
linear_module_scripted = torch.jit.script(linear_module)
model_ft_scripted_loaded.fc = linear_module_scripted
But this time the error is:
Expected a value of type '__torch__.torch.nn.modules.linear.Linear (of Python compilation unit at: 0x7f6bf29a53d0)' for field 'fc', but found '__torch__.torch.nn.modules.linear.___torch_mangle_25.Linear (of Python compilation unit at: 0x7f6bf31ff110)'
I searched online the the forum but I don’t find too much information to get me through this.
Any suggestions are appreciated!
Also, I wonder if there is a way to copy the weights from loaded scripted model. In that case, if I have the model (say, Resnet18) then I could copy the weights to the model and replace the last layer without problems.
P.S. My use case is to save and load models in different machines. But both in python. Thanks so much!