Pickle error when try to save a trained mode

1

I use pycox repository on Github for my project that is about survival analysis. This repository is based on Pytorch and torch tuples. After training a model we can save the model with

model.save_net('mynet.pt')

But when I use this code I face Pickle error that is not usual for other people. This error:

File ~/anaconda3/envs/myenv/lib/python3.9/site-packages/torchtuples/base.py:681 in save_net
    return torch.save(self.net, path, **kwargs)

  File ~/anaconda3/envs/myenv/lib/python3.9/site-packages/torch/serialization.py:380 in save
    _save(obj, opened_zipfile, pickle_module, pickle_protocol)

  File ~/anaconda3/envs/myenv/lib/python3.9/site-packages/torch/serialization.py:589 in _save
    pickler.dump(obj)

AttributeError: Can't pickle local object

Can anyone help me what could be the reason for this error?

2

I don’t know how save_net is implemented, but I would recommend to store the state_dict of the model instead of trying to pickle the model directly as it can fail in multiple ways during the loading (e.g. if you’ve changed the file structure or any function/class definition).

1 Like
3

Because I use a repository and I do not save the model myself it is hard for me to understand how can I save it with state_dict. This is the save_net part of code in repository:

 def save_net(self, path, **kwargs):
        """Save self.net and baseline hazards to file.
        Arguments:
            path {str} -- Path to file.
            **kwargs are passed to torch.save
        Returns:
            None
        """
        path, extension = os.path.splitext(path)
        if extension == "":
            extension = '.pt'
        super().save_net(path+extension, **kwargs)
        if hasattr(self, 'baseline_hazards_'):
            self.baseline_hazards_.to_pickle(path+'_blh.pickle')
4

Hi @maryma,

I am working with the DeepHit model in the pycox package.

Lets say that you trained a model “M”. You can save the trained weights as follows -

M.save_model_weights(PATH)

Here, “PATH” is the location where you want to save the weights.

To load these weights back into an empty model (say M*), you can just do -

M*.load_model_weights(PATH)

Both these methods are part of the DeepHit class. In Deep Learning, it is a common practice to not save the entire model architecture but just the weights, which makes deploying these models very efficient.

Hope this helps.
Ani