浅谈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中的CURL PycURL使用例子
Jun 01 Python
Python中列表(list)操作方法汇总
Aug 18 Python
探寻python多线程ctrl+c退出问题解决方案
Oct 23 Python
Python 装饰器使用详解
Jul 29 Python
Python实现PS图像明亮度调整效果示例
Jan 23 Python
python实现k-means聚类算法
Feb 23 Python
Linux下Pycharm、Anaconda环境配置及使用踩坑
Dec 19 Python
Python3.0中普通方法、类方法和静态方法的比较
May 03 Python
python 计算平均平方误差(MSE)的实例
Jun 29 Python
使用 Python ssh 远程登陆服务器的最佳方案
Mar 06 Python
python脚本和网页有何区别
Jul 02 Python
使用Python中tkinter库简单gui界面制作及打包成exe的操作方法(二)
Oct 12 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 中文和编码判断代码
2010/05/16 PHP
laravel安装zend opcache加速器教程
2015/03/02 PHP
php控制文件下载速度的方法
2015/03/24 PHP
PHP防盗链的基本思想 防盗链的设置方法
2015/09/25 PHP
Zend Framework实现多服务器共享SESSION数据的方法
2016/03/22 PHP
Eval and new funciton not the same thing
2012/12/27 Javascript
js 固定悬浮效果实现思路代码
2013/08/02 Javascript
NodeJS与Mysql的交互示例代码
2013/08/18 NodeJs
深入剖析JavaScript中的枚举功能
2014/03/06 Javascript
jQuery如何防止这种冒泡事件发生
2015/02/27 Javascript
jquery-tips悬浮提示插件分享
2015/07/31 Javascript
JavaScript判断用户名和密码不能为空的实现代码
2016/05/16 Javascript
javascript验证内容为数字以及长度为10的简单实例
2016/08/20 Javascript
VUE实现日历组件功能
2017/03/13 Javascript
VueJs 搭建Axios接口请求工具
2017/11/20 Javascript
让webpack+vue-cil项目不再自动打开浏览器的方法
2018/09/27 Javascript
Vue scoped及deep使用方法解析
2020/08/01 Javascript
Vue将props值实时传递 并可修改的操作
2020/08/09 Javascript
如何使用 vue-cli 创建模板项目
2020/11/19 Vue.js
[06:04]DOTA2英雄梦之声Vol19卓尔游侠
2014/06/20 DOTA
[00:36]我的中国心——Serenity vs Fnatic
2018/08/21 DOTA
python+selenium开发环境搭建图文教程
2017/08/11 Python
Python 实现文件打包、上传与校验的方法
2019/02/13 Python
Python基于Socket实现简易多人聊天室的示例代码
2020/11/29 Python
伦敦一家西班牙童装精品店:La Coqueta
2018/02/02 全球购物
财务出纳岗位职责
2014/02/03 职场文书
会计专业职业规划:规划自我赢取未来
2014/02/12 职场文书
物理研修随笔感言
2014/02/14 职场文书
体育教师个人工作总结
2015/02/09 职场文书
简单的辞职信模板
2015/05/12 职场文书
假期读书倡议书3篇
2019/08/19 职场文书
如何使用Python对NetCDF数据做空间相关分析
2021/04/21 Python
MySQL 数据类型选择原则
2021/05/27 MySQL
CSS+HTML 实现顶部导航栏功能
2021/08/30 HTML / CSS
python lambda 表达式形式分析
2022/04/03 Python
Linux服务器离线安装 nginx的详细步骤
2022/06/16 Servers