浅谈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 list使用示例 list中找连续的数字
Jan 27 Python
详解python如何调用C/C++底层库与互相传值
Aug 10 Python
听歌识曲--用python实现一个音乐检索器的功能
Nov 15 Python
Python 读取指定文件夹下的所有图像方法
Apr 27 Python
python2.x实现人民币转大写人民币
Jun 20 Python
PyQT实现菜单中的复制,全选和清空的功能的方法
Jun 17 Python
Python3远程监控程序的实现方法
Jul 15 Python
python支持多线程的爬虫实例
Dec 21 Python
Python高并发和多线程有什么关系
Nov 14 Python
python中os.remove()用法及注意事项
Jan 31 Python
使用Selenium实现微博爬虫(预登录、展开全文、翻页)
Apr 13 Python
浅谈Python numpy创建空数组的问题
May 25 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
PHP中判断变量为空的几种方法分享
2013/08/26 PHP
使用JavaScript创建新样式表和新样式规则
2016/06/14 PHP
从ThinkPHP3.2.3过渡到ThinkPHP5.0学习笔记图文详解
2019/04/03 PHP
javascript sudoku 数独智力游戏生成代码
2010/03/27 Javascript
javascript设计模式之工厂模式示例讲解
2014/03/04 Javascript
基于ajax实现文件上传并显示进度条
2015/08/03 Javascript
在web中js实现类似excel的表格控件
2016/09/01 Javascript
jQuery实现导航滚动到指定内容效果完整实例【附demo源码下载】
2016/09/20 Javascript
深入理解JS中的Function.prototype.bind()方法
2016/10/11 Javascript
老生常谈原生JS执行环境与作用域
2016/11/22 Javascript
为输入框加入数字js校验代码分享
2017/11/02 Javascript
Webpack优化配置缩小文件搜索范围
2017/12/25 Javascript
关于vue的语法规则检测报错问题的解决
2018/05/21 Javascript
vue-cli 3 全局过滤器的实例代码详解
2019/06/03 Javascript
node中使用log4js4.x版本记录日志的方法
2019/08/20 Javascript
小程序如何支持使用 async/await详解
2019/09/12 Javascript
微信小程序wxs实现吸顶效果
2020/01/08 Javascript
vue 导航守卫和axios拦截器有哪些区别
2020/12/19 Vue.js
使用Python读取安卓手机的屏幕分辨率方法
2018/03/31 Python
利用python GDAL库读写geotiff格式的遥感影像方法
2018/11/29 Python
Python3 SSH远程连接服务器的方法示例
2018/12/29 Python
利用ctypes获取numpy数组的指针方法
2019/02/12 Python
python实现简单银行管理系统
2019/10/25 Python
Python如何使用BeautifulSoup爬取网页信息
2019/11/26 Python
Python字典底层实现原理详解
2019/12/18 Python
Python 从attribute到property详解
2020/03/05 Python
python中怎么表示空值
2020/06/19 Python
浅谈Python3中print函数的换行
2020/08/05 Python
通过代码实例解析Pytest运行流程
2020/08/20 Python
幼儿园春季开学寄语
2014/04/03 职场文书
婚礼答谢词
2015/01/04 职场文书
2015年学雷锋活动总结
2015/02/06 职场文书
2015年法院工作总结范文
2015/04/28 职场文书
校友会致辞
2015/07/30 职场文书
病房管理制度范本
2015/08/06 职场文书
详解Redis基本命令与使用场景
2021/06/01 Redis