Skip to content

SAVE AND LOAD THE MODEL

There are two common uses for loading and saving models:

  • Save the model that has been trained to continue training next time.
  • Save the trained model for direct prediction in the future.

We will introduce how to use save and load to save and load models as follows.

Also, we will show how to load a pre-trained model for inference.

Saving and Loading Model Parameters

Module provided by OneFlow and defined by users provides the state_dict method to obtain all the model parameters, which is stored in a dictionary with the format "name-value".

import oneflow as flow
m = flow.nn.Linear(2,3)
print(m.state_dict())

The above code first constructs a Linear object, then prints its parameters.

OrderedDict([('weight',
              tensor([[-0.4297, -0.3571],
                      [ 0.6797, -0.5295],
                      [ 0.4918, -0.3039]], dtype=oneflow.float32, requires_grad=True)),
             ('bias',
              tensor([ 0.0977,  0.1219, -0.5372], dtype=oneflow.float32, requires_grad=True))])

We can load parameters by calling load_state_dict method of Module, as the following code:

myparams = {"weight":flow.ones(3,2), "bias":flow.zeros(3)}
m.load_state_dict(myparams)
print(m.state_dict())

The tensor in the dictionary created by us has been loaded into m Module:

OrderedDict([('weight',
              tensor([[1., 1.],
                      [1., 1.],
                      [1., 1.]], dtype=oneflow.float32, requires_grad=True)),
             ('bias',
              tensor([0., 0., 0.], dtype=oneflow.float32, requires_grad=True))])

Saving Models

We can use oneflow.save to save models.

flow.save(m.state_dict(), "./model")

The first parameter is the Module parameters, and the second is the saved path. The above code saves the parameters of the m Module object to the path ./model.

Loading Models

Using oneflow.load to load parameters from disk to memory with the specified path, and get the dictionary of the parameters.

params = flow.load("./model")

Then use load_state_dict to load the dictionary into the model.

m2 = flow.nn.Linear(2,3)
m2.load_state_dict(params)
print(m2.state_dict())

We have created a new Linear Module object m2, and loaded the parameters saved from the above to m2. Then we get the output as below:

OrderedDict([('weight', tensor([[1., 1.],
        [1., 1.],
        [1., 1.]], dtype=oneflow.float32, requires_grad=True)), ('bias', tensor([0., 0., 0.], dtype=oneflow.float32, requires_grad=True))])

Using a Pre-trained Model to Make Predictions

OneFlow can directly load PyTorch's pre-trained model for prediction as long as the structure and parameter names of the model are aligned with the PyTorch model.

Examples can be found in here.

Run commands below for trying how to use the pre-trained model to make predictions:

git clone https://github.com/Oneflow-Inc/models.git
cd models/Vision/classification/image/shufflenetv2/
bash infer.sh
Back to top