TensorFlow模型保存与重载实战：基于rasbt/deeplearning-models的MLP实现-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00860/article/details/148375630

TensorFlow模型保存与重载实战：基于rasbt/deeplearning-models的MLP实现

deeplearning-models A collection of various deep learning architectures, models, and tips 项目地址: https://gitcode.com/gh_mirrors/de/deeplearning-models

引言

在深度学习项目开发过程中，模型的保存与重载是至关重要的环节。本文将基于rasbt/deeplearning-models项目中的多层感知机(MLP)实现，深入讲解TensorFlow中模型保存与重载的多种策略。

模型保存的基本概念

在TensorFlow中，模型保存主要包含两部分内容：

模型结构：通过.meta文件保存计算图结构
模型参数：通过检查点(checkpoint)文件保存训练得到的权重和偏置

多层感知机(MLP)实现

我们首先定义一个具有两个隐藏层的MLP模型：

def fc_layer(input_tensor, n_output_units, name,
             activation_fn=None, seed=None,
             weight_params=None, bias_params=None):
    # 全连接层实现
    pass

def mlp_graph(n_input=784, n_classes=10, n_hidden_1=128, n_hidden_2=256,
              learning_rate=0.1,
              fixed_params=None):
    # 完整的MLP图结构定义
    pass

这个实现特别之处在于支持从NumPy数组直接加载参数，从而可以创建不可训练的"冻结"模型。

模型训练与保存

训练过程采用MNIST数据集，训练完成后使用tf.train.Saver保存模型：

with tf.Session(graph=g) as sess:
    # 训练过程...
    saver0 = tf.train.Saver()
    saver0.save(sess, save_path='./mlp')

这会生成三个文件：

.meta文件：保存模型结构
.data文件：保存模型参数
.index文件：保存变量索引

模型重载方法一：从检查点文件恢复

这是最常用的模型恢复方式：

with tf.Session() as sess:
    saver1 = tf.train.import_meta_graph('./mlp.meta')
    saver1.restore(sess, save_path='./mlp')
    # 使用恢复的模型进行预测...

这种方式完整恢复了模型结构和参数，可以继续训练或进行推理。

模型重载方法二：使用NumPy存档文件

有时我们需要将模型参数导出为更通用的格式，可以使用NumPy的.npz文件：

导出参数到NPZ文件：

params = {}
for v in var_names:
    ary = sess.run(v)
    params[v] = ary
np.savez('mlp', **params)

从NPZ文件加载创建不可训练模型：

param_dict = np.load('mlp.npz')
mlp_graph(fixed_params=param_dict)

这种方式的特别之处在于创建的是不可训练模型，适用于生产环境中的纯推理场景。如需继续训练，只需简单修改fc_layer函数，将tf.constant替换为tf.Variable。

两种方式的对比

| 特性 | 检查点文件方式 | NumPy存档方式 | |---------------------|--------------|--------------| | 保存模型结构 | 是 | 否 | | 保存模型参数 | 是 | 是 | | 恢复后可继续训练 | 是 | 需特殊处理 | | 跨平台兼容性 | 一般 | 优秀 | | 文件大小 | 较大 | 较小 |