PyTorch and Pickle Error: AttributeError: Can't get attribute "Class_name" on <module '__main__' from '/load_model.py'>

Question

I have created a neural network model using PyTorch for a time-series forecasting problem. I have saved the model using Pickle. When loading the model to check on test data, it is throwing an attribute error. Below, I'm providing the relevant parts of my code:

import torch
import torch.nn as nn
import torch.nn.functional as F

class twoD_predict(nn.Module):

  def __init__(self):

    super().__init__()

  def forward(self,x):
  ...

  def train(self,
        epochs = 100):
  ...

obj1 = twoD_predict()

obj1.train()

I saved the model using Pickle as follows:

import pickle
filename = (f"{column_names[0]}.sav")
pickle.dump(obj1, open(filename, 'wb'))

However, when I try to load the model with the following code:

import pickle
if __name__ == '__main__':
    with open("model.sav", 'rb') as file:
        model = pickle.load(file)

I encounter the error: "AttributeError: Can't get attribute "twoD_predict" on <module '__main__' from '/load_model.py'>"

Can anyone please help?

Answer 1

Have you try import twoD_predict to the file you want to load the model? The pickled version still requires things to replicate on how to recreate the structure of your model (e.g. you have to load the definition of the model to your current file). If you don't want to do that, considering export your model in TorchScript format.