Gensim's Doc2Vec/FastText/Word2Vec models with `callbacks` can't be loaded after saving

[darentsia]

Description

Gensim's Doc2Vec model can not be loaded after saving.
I trained my Doc2Vec model using examples of code provided below, saved it on 9th epoch and than when I'm trying to load it I receive an error (listing added).
No problems were received during training.
Training was executed on the same machine, on which I'm trying to load it again.

And please, provide the way to save this model to continue using it after training.

Steps/Code/Corpus to Reproduce

Training phase

class EpochSaver(CallbackAny2Vec):
    '''Callback to save model after each epoch and show training parameters '''

    def __init__(self, savedir):
        self.savedir = savedir
        self.epoch = 0
        os.makedirs(self.savedir, exist_ok=True)

    def on_epoch_end(self, model):
        savepath = os.path.join(self.savedir, "model_neg{}_epoch.gz".format(self.epoch))
        model.save(savepath)
        print(
            "Epoch saved: {}".format(self.epoch + 1),
            "Start next epoch ... ", sep="\n"
            )
        if os.path.isfile(os.path.join(self.savedir, "model_neg{}_epoch.gz".format(self.epoch - 1))):
            print("Previous model deleted ")
            os.remove(os.path.join(self.savedir, "model_neg{}_epoch.gz".format(self.epoch - 1))) 
        self.epoch += 1


def train():

    workers = multiprocessing.cpu_count()/2
    model = Doc2Vec(
        DocIter(),
        vec_size=700, alpha=0.03, min_alpha=0.00025, epochs=10,
        min_count=10, dm=1, hs=0, negative=10, workers=workers,
        window=20, callbacks=[EpochSaver("./checkpoints")]
    )  

Load trained model:

    from gensim.models.doc2vec import Doc2Vec
    model = Doc2Vec.load("checkpoints/model_neg9_epoch.gz")

Results

Loading model
INFO:gensim.utils:loading Doc2Vec object from checkpoints/model_neg9_epoch.gz
INFO:gensim.models.doc2vec:Model saved using code from earlier Gensim Version. Re-loading old model in a compatible way.
INFO:gensim.models.deprecated.old_saveload:loading Doc2Vec object from checkpoints/model_neg9_epoch.gz
Traceback (most recent call last):
File "/Users/user/Python/dynamic_topics/env/lib/python3.5/site-packages/gensim/models/doc2vec.py", line 689, in load
return super(Doc2Vec, cls).load(*args, **kwargs)
File "/Users/user/Python/dynamic_topics/env/lib/python3.5/site-packages/gensim/models/base_any2vec.py", line 629, in load
model = super(BaseWordEmbeddingsModel, cls).load(*args, **kwargs)
File "/Users/user/Python/dynamic_topics/env/lib/python3.5/site-packages/gensim/models/base_any2vec.py", line 278, in load
return super(BaseAny2VecModel, cls).load(fname_or_handle, **kwargs)
File "/Users/user/Python/dynamic_topics/env/lib/python3.5/site-packages/gensim/utils.py", line 425, in load
obj = unpickle(fname)
File "/Users/user/Python/dynamic_topics/env/lib/python3.5/site-packages/gensim/utils.py", line 1332, in unpickle
return _pickle.load(f, encoding='latin1')
AttributeError: Can't get attribute 'EpochSaver' on <module 'main' from 'cluster_d2v.py'>

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
File "cluster_d2v.py", line 71, in
model = Doc2Vec.load("checkpoints/model_neg9_epoch.gz")
File "/Users/user/Python/dynamic_topics/env/lib/python3.5/site-packages/gensim/models/doc2vec.py", line 693, in load
return load_old_doc2vec(*args, **kwargs)
File "/Users/user/Python/dynamic_topics/env/lib/python3.5/site-packages/gensim/models/deprecated/doc2vec.py", line 84, in load_old_doc2vec
old_model = Doc2Vec.load(*args, **kwargs)
File "/Users/user/Python/dynamic_topics/env/lib/python3.5/site-packages/gensim/models/deprecated/word2vec.py", line 1616, in load
model = super(Word2Vec, cls).load(*args, **kwargs)
File "/Users/user/Python/dynamic_topics/env/lib/python3.5/site-packages/gensim/models/deprecated/old_saveload.py", line 87, in load
obj = unpickle(fname)
File "/Users/user/Python/dynamic_topics/env/lib/python3.5/site-packages/gensim/models/deprecated/old_saveload.py", line 380, in unpickle
return _pickle.loads(file_bytes, encoding='latin1')
AttributeError: Can't get attribute 'EpochSaver' on <module 'main' from 'cluster_d2v.py'>

Versions

Darwin-17.6.0-x86_64-i386-64bit
Python 3.5.4 (v3.5.4:3f56838976, Aug 7 2017, 12:56:33)
[GCC 4.2.1 (Apple Inc. build 5666) (dot 3)]
NumPy 1.14.5
SciPy 1.1.0
gensim 3.4.0
FAST_VERSION 0

[piskvorky]

Thanks for reporting. That does look like a bug.

@gojomo @menshikh-iv @manneshiva any idea? To me it looks like the callback got serialized along with the model. Not sure if that's what we want.

@daridar can you try importing your EpochSaver from a module? (not just the main script, __main__).

[gojomo]

Not sure this is necessarily a functionality bug - that deserialization fails when a class is unavailable is a known limitation. May be a doc or error-catching insufficiency.

Perhaps the docs for callbacks property or CallbackAny2Vec superclass could be clearer. Perhaps risk could be caught/warned/errored on save() when effective classes look serialization-fragile. Perhaps error on load() could be clearer by catching/detecting the representative AttributeErrors and re-throwing a more descriptive error.

[darentsia]

Hello, @piskvorky and @gojomo ! Thank you for your kind advice.

I tried importing from file, where I created EpochSaver to file, where I'm trying to load my model and it helped!
Thank you, now everything works.

[piskvorky]

@daridar good to hear :) Thanks for confirming.

@gojomo @menshikh-iv Do we need to serialize the callbacks? Are they useful outside of training?

Now that I think about it, we probably do. Even here in @daridar 's use-case, if we don't serialize callbacks, a loaded "checkpoint" model wouldn't support any more check-pointing, which is potentially surprising (not good).

Let me close here as "won't fix" but a general documentation / error message update sounds like a good idea. @daridar what kind of documentation would have helped you? For example, did you read the documentation on models.Callbacks in full, would a clear warning there, somewhere near the top, have saved you the trouble? Or a better exception message?

[darentsia]

@piskvorky I think better exception message would have helped!

Because for me It was not so obvious that EpochSaver() needs to be imported from the module, where It was created while seeing this error) And, of course, It would be better if documentation would include some information about the correct usage of the saved model when using callbacks.

Thank you!

[piskvorky]

No problem. Can you open a PR that improves the docs?

Plus, if you're feeling a little more adventurous, a PR for the better exception message too :)

[menshikh-iv]

Do we need to serialize the callbacks?

IMO yes, we need (for the same reason as you suggest), update documentation with this note about serialization is a good idea I think.

[darentsia]

@piskvorky Hello! Was it addressed to me?

No problem. Can you open a PR that improves the docs?

Plus, if you're feeling a little more adventurous, a PR for the better exception message too :)

If it's for me, sorry, I missed it.

[piskvorky]

Yes, for you @daridar :)

Docstring / message fixes may seem insignificant, but they actually improve Gensim from an important perspective we (core devs) don't have -- the "newcomer user" perspective.

[JensMadsen]

I would prefer if it was possible to load models without loading classes/modules used to save the model in the first place. I know that would be a major refactoring. Anyhow, 1) I do not like to import modules that are not used, 2) I serve my prediction step on an AWS lambda (in production and for tests) and I train my model on an ec2 instance and therefore save my model in an s3 bucket. It is already problematic top fit the gensim doc2vec implementation in an AWs Lambda function mainly due lambda function size limitation (scipy is the big problem not gensim itself), so importing unnecessary modules makes my life more complicated because the import increase the lambda function size unless I refactor :-)

This is not a complaint; I really appreciate and enjoy using Gensim. Best Jens

[menshikh-iv]

Hi @JensMadsen, what exactly you mean by load models without loading classes/modules used to save the model in the first place?

If you need to load something, in gensim typically this looks like

from gensim.models import X

my_model = X.load(...)` 

my_model instance of class X -> X directly used here.

Please describe your idea in more details, we are open to discussion :)

[JensMadsen]

@menshikh-iv Sorry about being unclear :-)

Essentially my code is:

train_model.py:

def train_model(documents,
                args):   
    assert gensim.models.doc2vec.FAST_VERSION > -1

    if args.workers is None:
        args.workers = multiprocessing.cpu_count()
    logger.info('using {} workers'.format(args.workers))
    args_as_dict = vars(args)
    epoch_saver = epoch_saver.EpochSaver(args)
    model = gensim.models.doc2vec.Doc2Vec(documents=documents,
                                          callbacks=[epoch_saver],
                                          **args_as_dict)

    return model

epoch_saver.py:

class EpochSaver(gensim.models.callbacks.CallbackAny2Vec):

    def __init__(self, args):
        self.save_every_n_epoch = args.save_every_n_epoch
        self.args = args
        self.epoch = 0

    def on_epoch_end(self, model):
        if self.save_every_n_epoch is None:
            pass
        elif (self.epoch + 1) % self.save_every_n_epoch == 0:
            // save and upload to s3
            save_trained_model(model, self.args, epoch=self.epoch)
        self.epoch += 1

In order to load the model pickle need to have access to epoch_saver.py but epoch_saver.py is otherwise not in use. To infer vectors I therefore have to import epoch_saver.py in infer_vector.py which looks something like

import epoch_saver

def infer_docs(input_string, model_path=None, inferred_docs=10, steps=10):
    // when saved using a callback object this bit fails if pickle has no access to the class i.e. i must import epoch_saver ?
    model = load_model(model_path)
    
    logging.info('do simple preprocessing of input string')
    processed_str = simple_preprocess(input_string, min_len=2, max_len=35)
    logging.info('input string\n{}\nprocessed string{}\n'.format(
        input_string, processed_str))
    logging.info('infer vectors of input text')
    inferred_vector = model.infer_vector(processed_str, steps=steps)
    logging.info('calculating most similar docs')

    return model.docvecs.most_similar([inferred_vector], topn=inferred_docs)

That was my only point; importing epoch_saver.py in infer_vector.py without really using it in infer_vector.py :-) Since doc2vec has something like 4 levels of inheritance and since pickle is such an nice tool I know that doing it in a different way is very cumbersome :-)

[menshikh-iv]

@JensMadsen aha, I understand what you mean, thanks for the clarification!

In this case, probably a good idea to save the model without callback always (for avoiding any pickle issues), @piskvorky @gojomo wdyt?