Tensorflow2.0 保存和加载模型的几种方法

零、综述

• 1. save/load weights

• 2. save/load entire model
• 3. saved_model

一、Save the weights

1.一次性保存所有参数

model.save_weights('./checkpoints/my_checkpoint') 

2.加载权重

注意，用该方法保存模型只保存了参数，文件较小，加载较快，但是测试/部署时需要重建搭建网络。

model = create_model() #定义网络框架
model.load_weights('./checkpoints/my_checkpoint') #加载训练好的权重

loss, acc = model.evaluate(test_images, test_labels)

network.save_weights('weights.ckpt') #保存权重
print('saved weights')
del network

network = Sequential([layers.Dense(256)...])#模型必须跟训练时参数一模一样
network.compile(optimizer=optimizer.Adam(lr=0.01),loss=tf.losses.CategoricalCrossentropy(from_logits=True),
                metrics=['accuracy'])
network.load_weights('weights.ckpt') #加载后可从检查点处继续训练
network.evaluate(ds_val)

二.Save the model

该方法把模型也保存了,文件较大,效率比较低。

#保存模型和参数
network.save('model.h5')
#删除模型和参数
del network
#重新加载模型和参数
network = tf.keras.models.load_model('model.h5')
network.evaluate(x_val, y_val)

#三、ONNX
保存为onnx,这是通用格式,python生成的可以用c++解析，一般python训练而用C++部署。
注意，ONNX可以转TensorRT，以部署到NVIDIA的嵌入式设备中.

tf.saved_model.save(m, '/tmp/saved_model/') #可以给其余语言使用的

imported = tf.saved_model.load(path) #直接Load
f = imported.signatures["serving_default"]