Programmiererforum

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

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!