欢迎您访问程序员文章站本站旨在为大家提供分享程序员计算机编程知识!
您现在的位置是: 首页

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

程序员文章站 2024-03-14 14:29:28
...

零、综述

    1. save/load weights
    1. save/load entire model
    1. 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"]
相关标签: 代码示例