Saving models and visualizing them in netron

[pabloinigoblasco]

I am trying to visualize the models in netron. In order to do that I tried to save the model into the disk after loading from checkpoint.

    model = __load_model(model_name, pretrained_path)
    torch.save(model, "noise-filter.pth")

However, I get a pickle/serialization exception because of lambda usage inside modules:

Can't pickle local object 'SpectrogramUnet.__init__.<locals>.<lambda>'

Then I found a solution using the "dill" python package that is able to handle the lambda issues:

    model = __load_model(model_name, pretrained_path)
    torch.save(model, "noise-filter.pth",  pickle_module=dill)
    model = torch.load("noise-filter.pth" , pickle_module=dill)

After this, the model continues working and filtering properly the noise from some input (I think everything is correct). Nonetheless the model cannot still be open with netron.

Does anybody know what could be the issue?

PS: I also tried to export the model to onnx and visualize it in netron but I still get some errors.

[AppleHolic]

@pabloinigoblasco I checked out that the model has lambda funcion 1 and lambda function 2. It makes that the model cannot be pickled directly. I think you can fix them to save the model, and then it can be pickled.

[AppleHolic]

issue close