megengine.load#

load(f, map_location=None, pickle_module=pickle)[source]#

Load an object saved with save from a file.

Parameters:

  • f – a string of file name or a text file object from which to load.

  • map_location – defines device mapping. See examples for usage.

  • pickle_module – the module to use for pickling.

Note

If you will call set_default_device, please do it before load.

If you are using MegEngine with different Python versions

Different Python version may use different DEFAULT/HIGHEST pickle protocol. If you want to load the saved object in another Python version, please make sure you have used the same protocol.

You can select to use pickle module directly

This interface is a wrapper of pickle.load. If you want to use pickle, See pickle for more information about how to set pickle_protocol:

Examples

This example shows how to load tenors to different devices:

import megengine as mge

# Load tensors to the same device as defined in model.pkl
mge.load('model.pkl')

# Load all tensors to gpu0.
mge.load('model.pkl', map_location='gpu0')

# Load all tensors originally on gpu0 to cpu0
mge.load('model.pkl', map_location={'gpu0':'cpu0'})

# Load all tensors to cpu0
mge.load('model.pkl', map_location=lambda dev: 'cpu0')

If you are using a lower version of Python (<3.8), you can use other pickle module like pickle5 to load object saved in pickle 5 protocol:

>>> import pickle5 as pickle

Or you can use pickle5 in this way (only used with this interface)：

import pickle5
import megengine

megengine.load(obj, pickle_module=pickle5)