fast-reid/fastreid/export/tensorflow_export.py

# encoding: utf-8
"""
@author:  liaoxingyu
@contact: sherlockliao01@gmail.com
"""

import warnings

warnings.filterwarnings('ignore')  # Ignore all the warning messages in this tutorial
from onnx_tf.backend import prepare

import tensorflow as tf
from PIL import Image
import torchvision.transforms as transforms

import onnx
import numpy as np
import torch
from torch.backends import cudnn
import io

cudnn.benchmark = True


def _export_via_onnx(model, inputs):
    from ipdb import set_trace;
    set_trace()

    def _check_val(module):
        assert not module.training

    model.apply(_check_val)

    # Export the model to ONNX
    with torch.no_grad():
        with io.BytesIO() as f:
            torch.onnx.export(
                model,
                inputs,
                f,
                # verbose=True,  # NOTE: uncomment this for debugging
                export_params=True,
            )
            onnx_model = onnx.load_from_string(f.getvalue())
    # torch.onnx.export(model,  # model being run
    #                   inputs,  # model input (or a tuple for multiple inputs)
    #                   "reid_test.onnx",  # where to save the model (can be a file or file-like object)
    # export_params=True,  # store the trained parameter weights inside the model file
    # opset_version=10,  # the ONNX version to export the model to
    # do_constant_folding=True,  # whether to execute constant folding for optimization
    # input_names=['input'],  # the model's input names
    # output_names=['output'],  # the model's output names
    # dynamic_axes={'input': {0: 'batch_size'},  # variable lenght axes
    #               'output': {0: 'batch_size'}})
    # )

    # Apply ONNX's Optimization
    # all_passes = optimizer.get_available_passes()
    # passes = ["fuse_bn_into_conv"]
    # assert all(p in all_passes for p in passes)
    # onnx_model = optimizer.optimize(onnx_model, passes)

    # Convert ONNX Model to Tensorflow Model
    tf_rep = prepare(onnx_model, strict=False)  # Import the ONNX model to Tensorflow
    print(tf_rep.inputs)  # Input nodes to the model
    print('-----')
    print(tf_rep.outputs)  # Output nodes from the model
    print('-----')
    # print(tf_rep.tensor_dict)  # All nodes in the model
    # """

    # install onnx-tensorflow from github，and tf_rep = prepare(onnx_model, strict=False)
    # Reference https://github.com/onnx/onnx-tensorflow/issues/167
    # tf_rep = prepare(onnx_model) # whthout strict=False leads to KeyError: 'pyfunc_0'

    # debug, here using the same input to check onnx and tf.
    # output_onnx_tf = tf_rep.run(to_numpy(img))
    # print('output_onnx_tf = {}'.format(output_onnx_tf))
    # onnx --> tf.graph.pb
    # tf_pb_path = 'reid_tf_graph.pb'
    # tf_rep.export_graph(tf_pb_path)

    return tf_rep


def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()


def _check_pytorch_tf_model(model: torch.nn.Module, tf_graph_path: str):
    img = Image.open("demo_imgs/dog.jpg")

    resize = transforms.Resize([384, 128])
    img = resize(img)

    to_tensor = transforms.ToTensor()
    img = to_tensor(img)
    img.unsqueeze_(0)
    torch_outs = model(img)

    with tf.Graph().as_default():
        graph_def = tf.GraphDef()
        with open(tf_graph_path, "rb") as f:
            graph_def.ParseFromString(f.read())
            tf.import_graph_def(graph_def, name="")
        with tf.Session() as sess:
            # init = tf.initialize_all_variables()
            # init = tf.global_variables_initializer()
            # sess.run(init)

            # print all ops, check input/output tensor name.
            # uncomment it if you donnot know io tensor names.
            '''
            print('-------------ops---------------------')
            op = sess.graph.get_operations()
            for m in op:
                try:
                    # if 'input' in m.values()[0].name:
                    #     print(m.values())
                    if m.values()[0].shape.as_list()[1] == 2048: #and (len(m.values()[0].shape.as_list()) == 4):
                        print(m.values())
                except:
                    pass
            print('-------------ops done.---------------------')
            '''
            input_x = sess.graph.get_tensor_by_name('input.1:0')  # input
            outputs = sess.graph.get_tensor_by_name('502:0')  # 5
            output_tf_pb = sess.run(outputs, feed_dict={input_x: to_numpy(img)})

    print('output_pytorch = {}'.format(to_numpy(torch_outs)))
    print('output_tf_pb = {}'.format(output_tf_pb))

    np.testing.assert_allclose(to_numpy(torch_outs), output_tf_pb, rtol=1e-03, atol=1e-05)
    print("Exported model has been tested with tensorflow runtime, and the result looks good!")


def export_tf_reid_model(model: torch.nn.Module, tensor_inputs: torch.Tensor, graph_save_path: str):
    """
    Export a reid model via ONNX.
    Arg:
        model: a tf_1.x-compatible version of detectron2 model, defined in caffe2_modeling.py
        tensor_inputs: a list of tensors that caffe2 model takes as input.
    """
    # model = copy.deepcopy(model)
    assert isinstance(model, torch.nn.Module)

    # Export via ONNX
    print("Exporting a {} model via ONNX ...".format(type(model).__name__))
    predict_net = _export_via_onnx(model, tensor_inputs)
    print("ONNX export Done.")

    print("Saving graph of ONNX exported model to {} ...".format(graph_save_path))
    predict_net.export_graph(graph_save_path)

    print("Checking if tf.pb is right")
    _check_pytorch_tf_model(model, graph_save_path)

# if __name__ == '__main__':
# args = default_argument_parser().parse_args()
# print("Command Line Args:", args)
# cfg = setup(args)
# cfg = cfg.defrost()
# cfg.MODEL.BACKBONE.NAME = "build_resnet_backbone"
# cfg.MODEL.BACKBONE.DEPTH = 50
# cfg.MODEL.BACKBONE.LAST_STRIDE = 1
# # If use IBN block in backbone
# cfg.MODEL.BACKBONE.WITH_IBN = True
#
# model = build_model(cfg)
# # model.load_params_wo_fc(torch.load('logs/bjstation/res50_baseline_v0.4/ckpts/model_epoch80.pth'))
# model.cuda()
# model.eval()
# dummy_inputs = torch.randn(1, 3, 256, 128)
# export_tf_reid_model(model, dummy_inputs, 'reid_tf.pb')

# inputs = torch.rand(1, 3, 384, 128).cuda()
#
# _export_via_onnx(model, inputs)
# onnx_model = onnx.load("reid_test.onnx")
# onnx.checker.check_model(onnx_model)
#
# from PIL import Image
# import torchvision.transforms as transforms
#
# img = Image.open("demo_imgs/dog.jpg")
#
# resize = transforms.Resize([384, 128])
# img = resize(img)
#
# to_tensor = transforms.ToTensor()
# img = to_tensor(img)
# img.unsqueeze_(0)
# img = img.cuda()
#
# with torch.no_grad():
#     torch_out = model(img)
#
# ort_session = onnxruntime.InferenceSession("reid_test.onnx")
#
# # compute ONNX Runtime output prediction
# ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(img)}
# ort_outs = ort_session.run(None, ort_inputs)
# img_out_y = ort_outs[0]
#
#
# # compare ONNX Runtime and PyTorch results
# np.testing.assert_allclose(to_numpy(torch_out), ort_outs[0], rtol=1e-03, atol=1e-05)
#
# print("Exported model has been tested with ONNXRuntime, and the result looks good!")

# img = Image.open("demo_imgs/dog.jpg")
#
# resize = transforms.Resize([384, 128])
# img = resize(img)
#
# to_tensor = transforms.ToTensor()
# img = to_tensor(img)
# img.unsqueeze_(0)
# img = torch.cat([img.clone(), img.clone()], dim=0)

# ort_session = onnxruntime.InferenceSession("reid_test.onnx")

# # compute ONNX Runtime output prediction
# ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(img)}
# ort_outs = ort_session.run(None, ort_inputs)

# model = onnx.load('reid_test.onnx')  # Load the ONNX file
# tf_rep = prepare(model, strict=False) # Import the ONNX model to Tensorflow
# print(tf_rep.inputs)  # Input nodes to the model
# print('-----')
# print(tf_rep.outputs)  # Output nodes from the model
# print('-----')
# # print(tf_rep.tensor_dict)  # All nodes in the model

# install onnx-tensorflow from github，and tf_rep = prepare(onnx_model, strict=False)
# Reference https://github.com/onnx/onnx-tensorflow/issues/167
# tf_rep = prepare(onnx_model) # whthout strict=False leads to KeyError: 'pyfunc_0'

# # debug, here using the same input to check onnx and tf.
# # output_onnx_tf = tf_rep.run(to_numpy(img))
# # print('output_onnx_tf = {}'.format(output_onnx_tf))
# # onnx --> tf.graph.pb
# tf_pb_path = 'reid_tf_graph.pb'
# tf_rep.export_graph(tf_pb_path)

# # step 3, check if tf.pb is right.
# with tf.Graph().as_default():
#     graph_def = tf.GraphDef()
#     with open(tf_pb_path, "rb") as f:
#         graph_def.ParseFromString(f.read())
#         tf.import_graph_def(graph_def, name="")
#     with tf.Session() as sess:
#         # init = tf.initialize_all_variables()
#         init = tf.global_variables_initializer()
#         # sess.run(init)

#         # print all ops, check input/output tensor name.
#         # uncomment it if you donnot know io tensor names.
#         '''
#         print('-------------ops---------------------')
#         op = sess.graph.get_operations()
#         for m in op:
#             try:
#                 # if 'input' in m.values()[0].name:
#                 #     print(m.values())
#                 if m.values()[0].shape.as_list()[1] == 2048: #and (len(m.values()[0].shape.as_list()) == 4):
#                     print(m.values())
#             except:
#                 pass
#         print('-------------ops done.---------------------')
#         '''
#         input_x = sess.graph.get_tensor_by_name('input.1:0')  # input
#         outputs = sess.graph.get_tensor_by_name('502:0')  # 5
#         output_tf_pb = sess.run(outputs, feed_dict={input_x: to_numpy(img)})
#         print('output_tf_pb = {}'.format(output_tf_pb))
# np.testing.assert_allclose(ort_outs[0], output_tf_pb, rtol=1e-03, atol=1e-05)

# with tf.Graph().as_default():
#     graph_def = tf.GraphDef()
#     with open(tf_pb_path, "rb") as f:
#         graph_def.ParseFromString(f.read())
#         tf.import_graph_def(graph_def, name="")
#     with tf.Session() as sess:
#         # init = tf.initialize_all_variables()
#         init = tf.global_variables_initializer()
#         # sess.run(init)
#
#         # print all ops, check input/output tensor name.
#         # uncomment it if you donnot know io tensor names.
#         '''
#         print('-------------ops---------------------')
#         op = sess.graph.get_operations()
#         for m in op:
#             try:
#                 # if 'input' in m.values()[0].name:
#                 #     print(m.values())
#                 if m.values()[0].shape.as_list()[1] == 2048: #and (len(m.values()[0].shape.as_list()) == 4):
#                     print(m.values())
#             except:
#                 pass
#         print('-------------ops done.---------------------')
#         '''
#         input_x = sess.graph.get_tensor_by_name('input.1:0')  # input
#         outputs = sess.graph.get_tensor_by_name('502:0')  # 5
#         output_tf_pb = sess.run(outputs, feed_dict={input_x: to_numpy(img)})
#         from ipdb import set_trace;
#
#         set_trace()
#         print('output_tf_pb = {}'.format(output_tf_pb))