PicklingError

Thank you for the question @peterkim95 !

To load a general checkpoint, PyTorch provides a built-in function load_state_dict that you can call as follow:

checkpoint = torch.load('tmp/checkpoint-1000')
model.load_state_dict(checkpoint['model_state_dict'])

In Transformers4rec, we additionally simplified the model saving in the transformers4rec Trainer class (here) with a builtin method _save_model_and_checkpoint where you can save the checkpoints but the model class as well. By doing so, you don't have to re-define the model class.

An example of usage would be:

import cloudpickle
#Train and save the model 
recys_trainer = tr.Trainer(model, ...)
...
recsys_trainer._save_model_and_checkpoint(save_model_class=True)

# Load the model class and its checkpoint
checkpoint_path = 'tmp/checkpoint-1000'
model = cloudpickle.load(open(os.path.join(checkpoint_path, "model_class.pkl"), "rb"))

# Restoring model weights
model.load_state_dict(torch.load(os.path.join(checkpoint_path, "pytorch_model.bin")))

Let us know if those examples help you with your use-case :)

Originally posted by @sararb in https://github.com/NVIDIA-Merlin/Transformers4Rec/issues/348#issuecomment-993618831

When i am try this code, getting this error: PicklingError: Cannot pickle a prepared model with automatic mixed precision, please unwrap the model with Accelerator.unwrap_model(model) before pickling it.

If i try recsys_trainer.accelerator.unwrap_model and save again it is saving but at this time i am getting model.forward() missing 1 required positional argument: 'inputs'

My main aim is saving model and using like recsys_trainer.predict() it in another platform without triton server and GPU.

Start with transformers4rec/torch/trainer.py and the _save_model_and_checkpoint entry point referenced in the report, then trace the Accelerator unwrap path and the shown cloudpickle and load_state_dict calls. Reproduce saving under automatic mixed precision and loading for prediction on a platform without GPU; done means the reported pickling and missing-input errors no longer occur.