【4】TensorFlow光速入门-保存模型及加载模型并使用

提示：扫一扫查出行
【扫一扫了解最新限行尾号】
复制提示

本文地址：https://www.cnblogs.com/tujia/p/13862360.html

系列文章：

【0】TensorFlow光速入门-序

【1】TensorFlow光速入门-tensorflow开发基本流程

【2】TensorFlow光速入门-数据预处理（得到数据集）

【3】TensorFlow光速入门-训练及评估

【4】TensorFlow光速入门-保存模型及加载模型并使用

【5】TensorFlow光速入门-图片分类完整代码

【6】TensorFlow光速入门-python模型转换为tfjs模型并使用

【7】TensorFlow光速入门-总结

一、保存模型

创建一个目录

!mkdir /tf/saved_model

注：jupyter 代码块前面加一个!号表示，这是shell命令，不是代码；

保存模型

model.save('/tf/saved_model/wnw')

保存模型的其他参数及操作，看这里 https://tensorflow.google.cn/api_docs/python/tf/keras/Model#save

二、加载模型

import tensorflow as tf
from tensorflow import keras
import numpy as np
from IPython import display
import random# 加载模型
model = keras.models.load_model('/tf/saved_model/wnw')
# 看一下模型的结构
model.summary()# 随便找点图片
all_image_paths = []
data_root = pathlib.Path('/tf/datasets/wnw')
for item in data_root.rglob('*.jpg'):
all_image_paths.append(str(item))
print(len(all_image_paths))# 随机选取一张图片
img_path = random.choice(all_image_paths)
print(img_path)# 把图片处理成需要的tensor
image = tf.io.read_file(img_path)
image = tf.image.decode_image(image, channels=1)
image = tf.image.resize(image, (100, 100))
image /= 255
print(image.shape)# 预测只支持批量操作，我们给单张图片再加一维
images = (np.expand_dims(image, 0))
print(images.shape)
# 预测
predictions = model.predict(images)# 打印结果
label_names = ['other', 'watch']
label = np.argmax(predictions[0])
print(label_names[label])# 把图片也打印出来，看一下预测效果对不对
display.display(display.Image(img_path, width=200, height=200))