How to load a trained TF1 protobuf model into TF2?

Question

Update: This is a bug in tensorflow. Track progress here.

I have created and trained a model using stable-baselines, which uses Tensorflow 1. Now I need to use this trained model in an environment where I only have access to Tensorflow 2 or PyTorch. I figured I would go with Tensorflow 2 as the documentation says I should be able to load models created with Tensorflow 1.

I can load the pb file without a problem in Tensorflow 1:

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)

However in Tensorflow 2 I get the following error:

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].

Model definition:

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)

Model saving in 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})

Fully reproducible code can be found in the following 2 google colab notebooks: Tensorflow 1 saving and loading Tensorflow 2 loading

Direct link to the saved model: model

Answer 1

You can use compatibility layer of TensorFlow.

All v1 functionality is available under tf.compat.v1 namespace.

I managed to load your model in TF 2.1 (nothing special about that version, I just have it locally):

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>})