Xla: Loading Pytorch XLA-trained models in normal Pytorch?

[brianhhu]

I'm saving my trained models with something like this:

torch.save(model_parallel._models[0].state_dict(), 'model.pt')

However, when I transfer these saved models locally to test on a "normal" Pytorch installation, loading the saved model state_dict results in the following error:

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
c:\Users\brian\Documents\micronet\evaluate.py in <module>
    108 
    109 if __name__ == '__main__':
--> 110     main()

c:\Users\brian\Documents\micronet\evaluate.py in main()
     45 
     46     model = WRN_McDonnell(20, 10, 100)
---> 47     state_dict = torch.load('./checkpoints/model.pt')
     48 
     49 

c:\Users\brian\Miniconda3\envs\pytorch\lib\site-packages\torch\serialization.py in load(f, map_location, pickle_module, **pickle_load_args)
    384         f = f.open('rb')
    385     try:
--> 386         return _load(f, map_location, pickle_module, **pickle_load_args)
    387     finally:
    388         if new_fd:

c:\Users\brian\Miniconda3\envs\pytorch\lib\site-packages\torch\serialization.py in _load(f, map_location, pickle_module, **pickle_load_args)
    571     unpickler = pickle_module.Unpickler(f, **pickle_load_args)
    572     unpickler.persistent_load = persistent_load
--> 573     result = unpickler.load()
    574 
    575     deserialized_storage_keys = pickle_module.load(f, **pickle_load_args)

AttributeError: Can't get attribute '_rebuild_xlatensor' on <module 'torch._utils' from 'c:\\Users\\brian\\Miniconda3\\envs\\pytorch\\lib\\site-packages\\torch\\_utils.py'>

I may be missing something obvious here, but is there a solution to this? Thanks!

[dlibenzi]

In order to have loading back from checkpoint work with XLA we had to patch that part of the pytorch code:

https://github.com/pytorch/xla/blob/master/torch_patches/X10-torch_save.diff

So the pickler is failing there.
We will take a look into this ...

[ailzhang]

Hi @brianhhu , this is a known problem that will be fixed by upstreaming the change to pytorch/pytorch. https://github.com/pytorch/pytorch/pull/25882 needs a bit more work before merging. I'd expect this should be fixed by end of next week.
In the meanwhile, you can workaround this issue by loading the checkpoint where you have torch_xla installed, map all xla tensor back to cpu and save a new checkpoint with cpu tensors. Sorry for the inconvenience, let me know if this works for you or not.

[ibeltagy]

here's an implementation of the workaround,

state_dict = model_parallel._models[0].state_dict()
for t_name in state_dict:
   t_val = state_dict[t_name]
   state_dict[t_name] = t_val.to('cpu')
torch.save(state_dict,'model.pt')

[brianhhu]

Great, I will give this a try- thanks!

[dlibenzi]

This might work as well:

torch.save(model_parallel.models[0].to('cpu').state_dict(), 'model.pt')

[ailzhang]

This is "resolved" in the latest pytorch nightly, with a few caveats mentioned here. https://github.com/pytorch/xla/blob/master/API_GUIDE.md#discrepancies-between-pytorchxla
We'll revisit our serialization logic try to make it fully match PT logic. Please feel free to reopen if any of the caveats are blockers to you. Thanks!

[nosound2]

I was unaware of this thread, and was saving models with the following and loading to cpu model

state_dict = model_parallel._models[0].state_dict()
torch.save(state_dict, PATH_WORK/'models'/model_file_name)

It worked a month ago, but now the loading part does nothing, but also an error is not given. I did not expect the loading fail without an error. The loading loads nothing now apparently because it is saved on device:xla-1 and loads to cpu.

With the solutions proposed above, moving to cpu before saving, it works fine.

[dlibenzi]

Yeah, we have a section of the doc explaining the issue:

https://github.com/pytorch/xla/blob/master/API_GUIDE.md#saving-and-loading-xla-tensors