0

I am trying to use pick to save and load my ML models but I get an error. Here is the simplify version of my code to save my model:

import pickle
def test(x,y):
    return x+y

filename = 'test.pkl'
pickle.dump(test, open(filename, 'wb'))

I can load the pickle file from the same notebook that I am creating it but if I close the notebook and try to load the pick in a new one with the below code:

import pickle
filename = 'test.pkl'
loaded_model = pickle.load(open(filename, 'rb'))

It gets me this error:

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[2], line 2
      1 filename = 'test.pkl'
----> 2 loaded_model = pickle.load(open(filename, 'rb'))

AttributeError: Can't get attribute 'test' on <module '__main__'>
Amir
  • 1,017
  • 4
  • 14
  • 32
  • Does this answer your question? [Is there an easy way to pickle a python function (or otherwise serialize its code)?](https://stackoverflow.com/questions/1253528/is-there-an-easy-way-to-pickle-a-python-function-or-otherwise-serialize-its-cod) – Pranav Hosangadi Jan 10 '23 at 19:26
  • That post suggest using other packages which are not compatible with some of the Python ML libraries os I am looking for a solution to be able to use pickle – Amir Jan 10 '23 at 19:28
  • https://stackoverflow.com/questions/11866944/how-to-pickle-functions-classes-defined-in-main-python – Pranav Hosangadi Jan 10 '23 at 19:36
  • So this link suggests evey time I want to load my script/model it will re-run the whole script. Then what is the advantage of pickling the model? – Amir Jan 10 '23 at 20:44

1 Answers1

0

I had exactly the same issue as you. To fix this you need to import the function test in the file that you open the pickle file. In your example if the test function is located in the file ml_model.py, you need to add line from ml_model import test in the second file.