[Python] How to wrap a torch.jit model inside a torch Module?

I'm trying to call a TorchScript model inside a torch.nn.Module but got an error related to pickle.

Here's the code to reproduce:

import torch
import torch.nn as nn

# A simple base model to create a ScriptModel
class ExampleModel(nn.Module):
def __init__(self, factor: int):
super(ExampleModel, self).__init__()
self.factor = factor

def forward(self, x):
return x * self.factor

# Define a wrapper model with a ModuleDict
class WrapperModel(nn.Module):
def __init__(self, path):
super(WrapperModel, self).__init__()
self.model = torch.jit.load(path)


def forward(self, name: str, x):
return self.model(x)


scripted_model = torch.jit.script(ExampleModel(2))
scripted_model.save("model.jit")


# Initialize the WrapperModel
wrapper = WrapperModel("model.jit")


And when I try to pickle wrapper with:

import pickle
pickle.dumps(wrapper)


I got error:

RuntimeError: Tried to serialize object __torch__.___torch_mangle_3.ExampleModel which does not have a __getstate__ method defined!


Is there a way to call TorchScript model so that it doesn't raise such error?

Continue reading...