stable baseline export tensorflow
# Saving model. model here is a stable_baselines model
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})
# Loading model into TF1
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']
# important step
init = tf.global_variables_initializer()
global_session.run(init)
# Calling the model and converting the result to original format
computation_graph = model_loaded(tf.convert_to_tensor(grid))['action']
with global_session.as_default():
result = computation_graph.eval()
col = result.argmax()