浅谈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中super()函数简介及用法分享
Jul 11 Python
利用python将json数据转换为csv格式的方法
Mar 22 Python
python 移动图片到另外一个文件夹的实例
Jan 10 Python
python实现多层感知器
Jan 18 Python
python命令行工具Click快速掌握
Jul 04 Python
python字符串分割及字符串的一些常规方法
Jul 24 Python
Python调用接口合并Excel表代码实例
Mar 31 Python
keras模型保存为tensorflow的二进制模型方式
May 25 Python
浅谈Python中文件夹和python package包的区别
Jun 01 Python
python实现暗通道去雾算法的示例
Sep 27 Python
Python关于OS文件目录处理的实例分享
May 23 Python
python中的None与NULL用法说明
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记录日志的实现代码
2011/08/08 PHP
Zend的MVC机制使用分析(二)
2013/05/02 PHP
php多种形式发送邮件(mail qmail邮件系统 phpmailer类)
2014/01/22 PHP
由php中字符offset特征造成的绕过漏洞详解
2017/07/07 PHP
轻量级 JS ToolTip提示效果
2010/07/20 Javascript
js获取height和width的方法说明
2013/01/06 Javascript
使用jquery+CSS实现控制打印样式
2014/12/31 Javascript
JavaScript使用yield模拟多线程的方法
2015/03/19 Javascript
最简单的JavaScript图片轮播代码(两种方法)
2015/12/18 Javascript
微信小程序开发之toast等弹框提示使用教程
2017/06/08 Javascript
nodejs接入阿里大鱼短信验证码的方法
2017/07/10 NodeJs
把JavaScript代码改成ES6语法不完全指南(分享)
2017/09/10 Javascript
Vue精简版风格概述
2018/01/30 Javascript
如何在JavaScript中谨慎使用代码注释
2019/06/21 Javascript
JavaScript实现简易计算器小功能
2020/10/22 Javascript
[02:16]DOTA2超级联赛专访Burning 逆袭需要抓住机会
2013/06/24 DOTA
python数据结构树和二叉树简介
2014/04/29 Python
进一步理解Python中的函数编程
2015/04/13 Python
python任务调度实例分析
2015/05/19 Python
python调用API实现智能回复机器人
2018/04/10 Python
使用django的ORM框架按月统计近一年内的数据方法
2019/07/18 Python
Django框架基础模板标签与filter使用方法详解
2019/07/23 Python
Django实现分页显示效果
2019/10/31 Python
Python CSV文件模块的使用案例分析
2019/12/21 Python
Python爬取数据并实现可视化代码解析
2020/08/12 Python
阿联酋航空假期:Emirates Holidays
2018/03/20 全球购物
实习生体会的自我评价范文
2013/11/28 职场文书
淘宝网店营销策划书
2014/01/11 职场文书
新教师工作感言
2014/02/16 职场文书
语文课外活动总结
2014/08/27 职场文书
2014法院干警廉洁警示教育思想汇报
2014/09/13 职场文书
实习计划书范文
2015/01/16 职场文书
2015年车间安全管理工作总结
2015/05/13 职场文书
2015年国庆晚会主持词
2015/07/01 职场文书
小公司融资,商业计划书的8切记
2019/07/15 职场文书
python 实现图片特效处理
2022/04/03 Python