paddle动态图模型保存与加载

#模型的保存和加载

import numpy as np
import paddle
import paddle.nn as nn
import paddle.optimizer as opt

BATCH_SIZE = 16
BATCH_NUM = 4
EPOCH_NUM = 4

IMAGE_SIZE = 784
CLASS_NUM = 10

final_checkpoint = dict()

# 定义一个随机数据集
class RandomDataset(paddle.io.Dataset):
    def __init__(self, num_samples):
        self.num_samples = num_samples

    def __getitem__(self, idx):
        image = np.random.random([IMAGE_SIZE]).astype('float32')
        label = np.random.randint(0, CLASS_NUM - 1, (1, )).astype('int64')
        return image, label

    def __len__(self):
        return self.num_samples

class LinearNet(nn.Layer):
    def __init__(self):
        super().__init__()
        self._linear = nn.Linear(IMAGE_SIZE, CLASS_NUM)

    def forward(self, x):
        return self._linear(x)

def train(layer, loader, loss_fn, opt):
    for epoch_id in range(EPOCH_NUM):
        for batch_id, (image, label) in enumerate(loader()):
            out = layer(image)
            loss = loss_fn(out, label)
            loss.backward()
            opt.step()
            opt.clear_grad()
            print("Epoch {} batch {}: loss = {}".format(
                epoch_id, batch_id, np.mean(float(loss))))
            print('epoch{}'.format(epoch_id),'优化器',opt.state_dict())
        # 最后一个epoch保存检查点checkpoint
        if epoch_id == EPOCH_NUM - 1:
            final_checkpoint["epoch"] = epoch_id
            final_checkpoint["loss"] =float(loss)

# 创建网络、loss和优化器
layer = LinearNet()
loss_fn = nn.CrossEntropyLoss()
adam = opt.Adam(learning_rate=0.001, parameters=layer.parameters())

# 创建用于载入数据的DataLoader
dataset = RandomDataset(BATCH_NUM * BATCH_SIZE)
loader = paddle.io.DataLoader(dataset,
    batch_size=BATCH_SIZE,
    shuffle=True,
    drop_last=True,
    num_workers=0)
print('模型',type(layer),'模型结构',layer,'layer.state_dict()',layer.state_dict())
print('优化器',adam.state_dict())
# 开始训练
train(layer, loader, loss_fn, adam)
print('保存的对象',final_checkpoint)

##保存模型参数

# 保存Layer参数
paddle.save(layer.state_dict(), "linear_net.pdparams")
# 保存优化器参数
paddle.save(adam.state_dict(), "adam.pdopt")
# 保存检查点checkpoint信息
paddle.save(final_checkpoint, "final_checkpoint.pkl")


##加载动态图模型

final_checkpoint=paddle.load('final_checkpoint.pkl')
print('final_checkpoint',final_checkpoint)
layer_state_dict=paddle.load('linear_net.pdparams')
print('layer_state_dict',layer_state_dict)
opt_state_dict=paddle.load('adam.pdopt')
print('adam state_dict',opt_state_dict)