Wie wickle ich ein Fackel.jit -Modell in ein Taschenmodul ein? ⇐ Python

[Anonymous]

Ich versuche, ein Torchscript -Modell in einer Fackel zu rufen. p>

Code: Select all

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")

Und wenn ich versuche, Wrapper mit:
zu wickeln

Code: Select all

import pickle
pickle.dumps(wrapper)
< /code>
Ich habe Fehler erhalten: < /p>
RuntimeError: Tried to serialize object __torch__.___torch_mangle_3.ExampleModel which does not have a __getstate__ method defined!

Gibt es eine Möglichkeit, das Torchscript -Modell aufzurufen, damit ein solcher Fehler nicht erhöht wird?