浅谈Tensorflow模型的保存与恢复加载


Posted in Python onApril 26, 2018

近期做了一些反垃圾的工作,除了使用常用的规则匹配过滤等手段,也采用了一些机器学习方法进行分类预测。我们使用TensorFlow进行模型的训练,训练好的模型需要保存,预测阶段我们需要将模型进行加载还原使用,这就涉及TensorFlow模型的保存与恢复加载。

总结一下Tensorflow常用的模型保存方式。

保存checkpoint模型文件(.ckpt)

首先,TensorFlow提供了一个非常方便的api,tf.train.Saver()来保存和还原一个机器学习模型。

模型保存

使用tf.train.Saver()来保存模型文件非常方便,下面是一个简单的例子:

import tensorflow as tf
import os

def save_model_ckpt(ckpt_file_path):
  x = tf.placeholder(tf.int32, name='x')
  y = tf.placeholder(tf.int32, name='y')
  b = tf.Variable(1, name='b')
  xy = tf.multiply(x, y)
  op = tf.add(xy, b, name='op_to_store')

  sess = tf.Session()
  sess.run(tf.global_variables_initializer())

  path = os.path.dirname(os.path.abspath(ckpt_file_path))
  if os.path.isdir(path) is False:
    os.makedirs(path)

  tf.train.Saver().save(sess, ckpt_file_path)
  
  # test
  feed_dict = {x: 2, y: 3}
  print(sess.run(op, feed_dict))

程序生成并保存四个文件(在版本0.11之前只会生成三个文件:checkpoint, model.ckpt, model.ckpt.meta)

  1. checkpoint 文本文件,记录了模型文件的路径信息列表
  2. model.ckpt.data-00000-of-00001 网络权重信息
  3. model.ckpt.index .data和.index这两个文件是二进制文件,保存了模型中的变量参数(权重)信息
  4. model.ckpt.meta 二进制文件,保存了模型的计算图结构信息(模型的网络结构)protobuf

以上是tf.train.Saver().save()的基本用法,save()方法还有很多可配置的参数:

tf.train.Saver().save(sess, ckpt_file_path, global_step=1000)

加上global_step参数代表在每1000次迭代后保存模型,会在模型文件后加上"-1000",model.ckpt-1000.index, model.ckpt-1000.meta, model.ckpt.data-1000-00000-of-00001

每1000次迭代保存一次模型,但是模型的结构信息文件不会变,就只用1000次迭代时保存一下,不用相应的每1000次保存一次,所以当我们不需要保存meta文件时,可以加上write_meta_graph=False参数,如下:

tf.train.Saver().save(sess, ckpt_file_path, global_step=1000, write_meta_graph=False)

如果想每两小时保存一次模型,并且只保存最新的4个模型,可以加上使用max_to_keep(默认值为5,如果想每训练一个epoch就保存一次,可以将其设置为None或0,但是没啥用不推荐), keep_checkpoint_every_n_hours参数,如下:

tf.train.Saver().save(sess, ckpt_file_path, max_to_keep=4, keep_checkpoint_every_n_hours=2)

同时在tf.train.Saver()类中,如果我们不指定任何信息,则会保存所有的参数信息,我们也可以指定部分想要保存的内容,例如只保存x, y参数(可传入参数list或dict):

tf.train.Saver([x, y]).save(sess, ckpt_file_path)

ps. 在模型训练过程中需要在保存后拿到的变量或参数名属性name不能丢,不然模型还原后不能通过get_tensor_by_name()获取。

模型加载还原

针对上面的模型保存例子,还原模型的过程如下:

import tensorflow as tf

def restore_model_ckpt(ckpt_file_path):
  sess = tf.Session()
  saver = tf.train.import_meta_graph('./ckpt/model.ckpt.meta') # 加载模型结构
  saver.restore(sess, tf.train.latest_checkpoint('./ckpt')) # 只需要指定目录就可以恢复所有变量信息

  # 直接获取保存的变量
  print(sess.run('b:0'))

  # 获取placeholder变量
  input_x = sess.graph.get_tensor_by_name('x:0')
  input_y = sess.graph.get_tensor_by_name('y:0')
  # 获取需要进行计算的operator
  op = sess.graph.get_tensor_by_name('op_to_store:0')

  # 加入新的操作
  add_on_op = tf.multiply(op, 2)

  ret = sess.run(add_on_op, {input_x: 5, input_y: 5})
  print(ret)

首先还原模型结构,然后还原变量(参数)信息,最后我们就可以获得已训练的模型中的各种信息了(保存的变量、placeholder变量、operator等),同时可以对获取的变量添加各种新的操作(见以上代码注释)。
并且,我们也可以加载部分模型,在此基础上加入其它操作,具体可以参考官方文档和demo。

针对ckpt模型文件的保存与还原,stackoverflow上有一个回答解释比较清晰,可以参考。

同时cv-tricks.com上面的TensorFlow模型保存与恢复的教程也非常好,可以参考。

以上就是本文的全部内容,希望对大家的学习有所帮助,也希望大家多多支持三水点靠木。

Python 相关文章推荐
使用Python的web.py框架实现类似Django的ORM查询的教程
May 02 Python
谈谈Python进行验证码识别的一些想法
Jan 25 Python
详解Python的hasattr() getattr() setattr() 函数使用方法
Jul 09 Python
Python Series从0开始索引的方法
Nov 06 Python
Python实现繁?转为简体的方法示例
Dec 18 Python
Python函数定义及传参方式详解(4种)
Mar 18 Python
如何用Python来理一理红楼梦里的那些关系
Aug 14 Python
windows 10 设定计划任务自动执行 python 脚本的方法
Sep 11 Python
Django Form and ModelForm的区别与使用
Dec 06 Python
python matplotlib画盒图、子图解决坐标轴标签重叠的问题
Jan 19 Python
jupyter notebook 多环境conda kernel配置方式
Apr 10 Python
Pymysql实现往表中插入数据过程解析
Jun 02 Python
Python实现爬取百度贴吧帖子所有楼层图片的爬虫示例
Apr 26 #Python
Python实现的计算器功能示例
Apr 26 #Python
python email smtplib模块发送邮件代码实例
Apr 26 #Python
Python利用正则表达式实现计算器算法思路解析
Apr 25 #Python
Python实现随机生成手机号及正则验证手机号的方法
Apr 25 #Python
Python实现按中文排序的方法示例
Apr 25 #Python
Python实现的基于优先等级分配糖果问题算法示例
Apr 25 #Python
You might like
留言板翻页的实现详解
2006/10/09 PHP
PHP编程最快明白(第一讲 软件环境和准备工作)
2010/10/25 PHP
thinkPHP3.1验证码的简单实现方法
2016/04/22 PHP
yii框架搜索分页modle写法
2016/12/19 PHP
PHP通过文件路径获取文件名的实例代码
2018/10/14 PHP
关于Javascript 的 prototype问题。
2007/01/03 Javascript
jquery实现图片等比例缩放以及max-width在ie中不兼容解决
2013/03/21 Javascript
vue.js中v-on:textInput无法执行事件问题的解决过程
2017/07/12 Javascript
vue axios整合使用全攻略
2018/05/24 Javascript
jQuery实现的简单获取索引功能示例
2018/06/04 jQuery
关于微信公众号开发无法支付的问题解决
2018/12/28 Javascript
微信小程序实现复选框效果
2018/12/28 Javascript
微信小程序如何访问公众号文章
2019/07/08 Javascript
ES5新增数组的实现方法
2020/05/12 Javascript
Vue自定义全局弹窗组件操作
2020/08/11 Javascript
vue-cli4使用全局less文件中的变量配置操作
2020/10/21 Javascript
python实现的各种排序算法代码
2013/03/04 Python
总结Python编程中函数的使用要点
2016/03/20 Python
centos6.4下python3.6.1安装教程
2017/07/21 Python
Python3之读取连接过的网络并定位的方法
2018/04/22 Python
基于python 爬虫爬到含空格的url的处理方法
2018/05/11 Python
用Q-learning算法实现自动走迷宫机器人的方法示例
2019/06/03 Python
virtualenv介绍及简明教程
2020/06/23 Python
django数据模型中null和blank的区别说明
2020/09/02 Python
教你如何用python操作摄像头以及对视频流的处理
2020/10/12 Python
详解anaconda安装步骤
2020/11/23 Python
蛋白质世界:Protein World
2017/11/23 全球购物
几道数据库的面试题或笔试题
2014/05/31 面试题
应届毕业生求职信范文分享
2013/12/26 职场文书
数控技术专科生自我评价
2014/01/08 职场文书
管理建议书范文
2014/05/13 职场文书
校园广播稿精选
2014/10/01 职场文书
信息合作协议书
2014/10/09 职场文书
2014年消防工作总结
2014/11/21 职场文书
大学生学期个人总结
2015/02/12 职场文书
教你如何使用Python Tkinter库制作记事本
2021/06/10 Python