如何将训练有素的TF1原型模型加载到TF2中？

编程问答 2022-08-13

问题描述

更新：这是张量流中的错误。跟踪进度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,padding='same')
CONV3 = Conv2D(NUM_CHANNELS,strides=1)
CONV4 = Conv2D(NUM_CHANNELS,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

直接链接到保存的模型： model

解决方法

您可以使用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>})
stable-baselinestensorflowtensorflowtensorflowtensorflow2.0