使用 tensorflow 导出神经网络的权重

Question

我使用 tensorflow 工具编写了神经网络。 一切正常，现在我想导出我的神经网络的最终权重来做一个单一的预测方法。 我怎样才能做到这一点？

Answer 1

您需要在训练结束时使用tf.train.Saver类保存模型。

在初始化Saver对象时，您需要传递您希望保存的所有变量的列表。 最好的部分是您可以在不同的计算图中使用这些保存的变量！

通过使用创建一个Saver对象，

# Assume you want to save 2 variables `v1` and `v2`
saver = tf.train.Saver([v1, v2])

使用tf.Session对象保存变量，

saver.save(sess, 'filename');

当然，您可以添加其他详细信息，例如global_step 。

您可以在将来使用restore()函数恢复变量。 恢复的变量将自动初始化为这些值。

Answer 2

上面的答案是保存/恢复会话快照的标准方法。 但是，如果要将网络导出为单个二进制文件以进一步与其他 tensorflow 工具一起使用，则需要执行更多步骤。

首先， 冻结图表。 TF 提供了相应的工具。 我这样使用它：

#!/bin/bash -x

# The script combines graph definition and trained weights into
# a single binary protobuf with constant holders for the weights.
# The resulting graph is suitable for the processing with other tools.


TF_HOME=~/tensorflow/

if [ $# -lt 4 ]; then
    echo "Usage: $0 graph_def snapshot output_nodes output.pb"
    exit 0
fi

proto=$1
snapshot=$2
out_nodes=$3
out=$4

$TF_HOME/bazel-bin/tensorflow/python/tools/freeze_graph --input_graph=$proto \
    --input_checkpoint=$snapshot \
    --output_graph=$out \
    --output_node_names=$out_nodes 

完成此操作后，您可以对其进行优化以进行推理，或使用任何其他工具。

Answer 3

如果您只需要访问神经网络的权重和偏差，您可以使用 tf.keras.layers.Layer 中的tf.keras.layers.Layer get_weights()方法

#params vector includes weights and biases
params = your_model.get_weights()