如何将经过训练的 TF1 protobuf 模型加载到 TF2 中？

Question

更新： 这是tensorflow的一个bug。跟踪进度 here.

我使用稳定基线创建并训练了一个模型，该模型使用 Tensorflow 1。现在我需要在只能访问 Tensorflow 2 或 PyTorch 的环境中使用这个经过训练的模型。我想我会选择 Tensorflow 2 作为 documentation says 我应该能够加载使用 Tensorflow 1 创建的模型。

我可以在 Tensorflow 1 中毫无问题地加载 pb 文件:

global_session = tf.Session()

with global_session.as_default():
    model_loaded = tf.saved_model.load_v2('tensorflow_model')
    model_loaded = model_loaded.signatures['serving_default']

init = tf.global_variables_initializer()
global_session.run(init)

但是在 Tensorflow 2 中我得到以下错误:

can_be_imported = tf.saved_model.contains_saved_model('tensorflow_model')
assert(can_be_imported)
model_loaded = tf.saved_model.load('tensorflow_model/')

ValueError: Node 'loss/gradients/model/batch_normalization_3/FusedBatchNormV3_1_grad/FusedBatchNormGradV3' has an _output_shapes attribute inconsistent with the GraphDef for output #3: Dimension 0 in both shapes must be equal, but are 0 and 64. Shapes are [0] and [64].

模型定义:

NUM_CHANNELS = 64

BN1 = BatchNormalization()
BN2 = BatchNormalization()
BN3 = BatchNormalization()
BN4 = BatchNormalization()
BN5 = BatchNormalization()
BN6 = BatchNormalization()
CONV1 = Conv2D(NUM_CHANNELS, kernel_size=3, strides=1, padding='same')
CONV2 = Conv2D(NUM_CHANNELS, kernel_size=3, strides=1, padding='same')
CONV3 = Conv2D(NUM_CHANNELS, kernel_size=3, strides=1)
CONV4 = Conv2D(NUM_CHANNELS, kernel_size=3, strides=1)
FC1 = Dense(128)
FC2 = Dense(64)
FC3 = Dense(7)

def modified_cnn(inputs, **kwargs):
    relu = tf.nn.relu
    log_softmax = tf.nn.log_softmax
    
    layer_1_out = relu(BN1(CONV1(inputs)))
    layer_2_out = relu(BN2(CONV2(layer_1_out)))
    layer_3_out = relu(BN3(CONV3(layer_2_out)))
    layer_4_out = relu(BN4(CONV4(layer_3_out)))
    
    flattened = tf.reshape(layer_4_out, [-1, NUM_CHANNELS * 3 * 2]) 
    
    layer_5_out = relu(BN5(FC1(flattened)))
    layer_6_out = relu(BN6(FC2(layer_5_out)))
    
    return log_softmax(FC3(layer_6_out))

class CustomCnnPolicy(CnnPolicy):
    def __init__(self, *args, **kwargs):
        super(CustomCnnPolicy, self).__init__(*args, **kwargs, cnn_extractor=modified_cnn)

model = PPO2(CustomCnnPolicy, env, verbose=1)

模型保存在 TF1 中：

with model.graph.as_default():
    tf.saved_model.simple_save(model.sess, 'tensorflow_model', inputs={"obs": model.act_model.obs_ph},
                                   outputs={"action": model.act_model._policy_proba})

可以在以下 2 个 google colab 笔记本中找到完全可重现的代码： Tensorflow 1 saving and loading Tensorflow 2 loading

直接link到保存的模型： model

Answer 1

您可以使用 TensorFlow 的兼容层。

所有 v1 功能在 tf.compat.v1 命名空间下可用。

我成功地在 TF 2.1 中加载了您的模型（那个版本没有什么特别的，我只是在本地有）：

import tensorflow as tf

tf.__version__
Out[2]: '2.1.0'

model = tf.compat.v1.saved_model.load_v2('~/tmp/tensorflow_model')

model.signatures
Out[3]: _SignatureMap({'serving_default': <tensorflow.python.eager.wrap_function.WrappedFunction object at 0x7ff9244a6908>})

如何将经过训练的 TF1 protobuf 模型加载到 TF2 中？

How to load a trained TF1 protobuf model into TF2?

tensorflow

tensorflow2.0

stable-baselines