How can I save / reload a FunctionTransformer object and expect it to always work the same way, even if its internal function definition changed?

Viewed 35

I am using ColumnTransformer and FunctionTransformer classes to build a preprocessing pipeline for an ML project.

One of my FunctionTransformer uses a function (let's say preprocess_A) defined in my package. I fitted the pipeline and saved it as a pickle file alongside a trained model. Everything seemed to work fine.

I then decided to change preprocess_A definition for a new experiment, but I also noticed a drop in performance for my trained model. The problem is that the pickle file only keeps a reference to preprocess_A but not its definition. Hence, any change in preprocess_A will impact my previous pipelines.

So, is there a way to save my FunctionTransformer object and expect it to always work the same way, even if I later modify preprocess_A?

Code snippet :

### Imports
import pickle
import pandas as pd
# Custom package pipeline
from my_package.preprocessing import preprocess
pipeline = preprocess.get_pipeline()
# pipeline = ColumnTransformer([('test', FunctionTransformer(preprocess_A), make_column_selector())])
# def preprocess_A(x): return x ** 2
### Fit
df = pd.DataFrame({'col1': [1, 2, 4], 'col2': [3, 6, 9]})
pipeline.fit(df)
### Save
with open('test.pkl', 'wb') as f:
    pickle.dump(pipeline, f)
### Transform
pipeline.transform(df)
# array([[ 1,  9], [ 4, 36], [16, 81]], dtype=int64)
- def preprocess_A(x): return x ** 2
+ def preprocess_A(x): return x ** 4
### Imports
import pickle
import pandas as pd
### Load pipeline
with open('test.pkl', 'rb') as f:
    pipeline = pickle.load(f)
### Transform
df = pd.DataFrame({'col1': [1, 2, 4], 'col2': [3, 6, 9]})
pipeline.transform(df)
# array([[   1,   81], [  16, 1296], [ 256, 6561]], dtype=int64)

Note : it works fine with dill and lambda functions, but some preprocessing function can be complex and not easy to transform in lambdas

Note 2 : even if some functions might be complex, they have no dependency to my own package

1 Answers

I'm the dill author. With pickle, the behavior is to load the class instance where the class instance refers to the class object by reference (i.e. so if the class definition changes, then the unserialized object uses the new version). dill on the other hand, has a keyword that allows you to toggle the behavior. dill stores the class definition along with the instance, so you can choose to either use or ignore the stored class definition.

Python 3.7.14 (default, Sep 10 2022, 11:17:06) 
[Clang 10.0.1 (clang-1001.0.46.4)] on darwin
Type "help", "copyright", "credits" or "license" for more information.
>>> 
>>> import dill
>>> class Foo(object):
...   def bar(self, x):
...     return x+self.y
...   y = 1
... 
>>> f = Foo()
>>> _Foo = dill.dumps(Foo)
>>> _f = dill.dumps(f)
>>> del Foo, f
>>> 
>>> class Foo(object):
...   def bar(self, x):
...     return x+self.z
...   z = -1
... 
>>> f_ = dill.loads(_f) # use the new class
>>> f_.__class__ == Foo
True
>>> f_.z
-1
>>> f_.bar(0)
-1
>>> f_ = dill.loads(_f, ignore=True) # ignore the new class
>>> f_.bar(0)
1
>>> f_.y
1
>>> 

EDIT:

If it's a imported class, instead of defined in __main__, it's a bit harder, but you still might be able to do it using dill.source.

>>> import dill
>>> from changeclass import Foo
>>> # stored by reference, so not what you want
>>> dill.dumps(Foo)
b'\x80\x03cchangeclass\nFoo\nq\x00.'
>>> # so, we grab the source
>>> dill.source.getsource(Foo)
'class Foo(object):\n  def bar(self, x):\n    return self.y + x\n  y = 1\n'
>>> 
>>> # we can either exec the source here and then
>>> # store an instance to it (see above answer),
>>> # or we can dump the source and exec it later.
>>> 
>>> dill.dumps(dill.source.getsource(Foo))
b'\x80\x03XE\x00\x00\x00class Foo(object):\n  def bar(self, x):\n    return self.y + x\n  y = 1\nq\x00.'
>>> dill.loads(_)
'class Foo(object):\n  def bar(self, x):\n    return self.y + x\n  y = 1\n'
>>> exec(_)
>>> Foo
<class '__main__.Foo'>

If the class is complicated, you might have to dump the entire module source... or worse, multiple modules. It can get complicated. It's not intended to work across different source versions.

Related