Arulkumar03/Wheat_HEAD_Detection_Counting_ComputerVision_Model

1# Copyright (c) Facebook, Inc. and its affiliates.2 3import copy4import io5import logging6import numpy as np7from typing import List8import onnx9import onnx.optimizer10import torch11from caffe2.proto import caffe2_pb212from caffe2.python import core13from caffe2.python.onnx.backend import Caffe2Backend14from tabulate import tabulate15from termcolor import colored16from torch.onnx import OperatorExportTypes17 18from .shared import (19    ScopedWS,20    construct_init_net_from_params,21    fuse_alias_placeholder,22    fuse_copy_between_cpu_and_gpu,23    get_params_from_init_net,24    group_norm_replace_aten_with_caffe2,25    infer_device_type,26    remove_dead_end_ops,27    remove_reshape_for_fc,28    save_graph,29)30 31logger = logging.getLogger(__name__)32 33 34def export_onnx_model(model, inputs):35    """36    Trace and export a model to onnx format.37 38    Args:39        model (nn.Module):40        inputs (tuple[args]): the model will be called by `model(*inputs)`41 42    Returns:43        an onnx model44    """45    assert isinstance(model, torch.nn.Module)46 47    # make sure all modules are in eval mode, onnx may change the training state48    # of the module if the states are not consistent49    def _check_eval(module):50        assert not module.training51 52    model.apply(_check_eval)53 54    # Export the model to ONNX55    with torch.no_grad():56        with io.BytesIO() as f:57            torch.onnx.export(58                model,59                inputs,60                f,61                operator_export_type=OperatorExportTypes.ONNX_ATEN_FALLBACK,62                # verbose=True,  # NOTE: uncomment this for debugging63                # export_params=True,64            )65            onnx_model = onnx.load_from_string(f.getvalue())66 67    return onnx_model68 69 70def _op_stats(net_def):71    type_count = {}72    for t in [op.type for op in net_def.op]:73        type_count[t] = type_count.get(t, 0) + 174    type_count_list = sorted(type_count.items(), key=lambda kv: kv[0])  # alphabet75    type_count_list = sorted(type_count_list, key=lambda kv: -kv[1])  # count76    return "\n".join("{:>4}x {}".format(count, name) for name, count in type_count_list)77 78 79def _assign_device_option(80    predict_net: caffe2_pb2.NetDef, init_net: caffe2_pb2.NetDef, tensor_inputs: List[torch.Tensor]81):82    """83    ONNX exported network doesn't have concept of device, assign necessary84    device option for each op in order to make it runable on GPU runtime.85    """86 87    def _get_device_type(torch_tensor):88        assert torch_tensor.device.type in ["cpu", "cuda"]89        assert torch_tensor.device.index == 090        return torch_tensor.device.type91 92    def _assign_op_device_option(net_proto, net_ssa, blob_device_types):93        for op, ssa_i in zip(net_proto.op, net_ssa):94            if op.type in ["CopyCPUToGPU", "CopyGPUToCPU"]:95                op.device_option.CopyFrom(core.DeviceOption(caffe2_pb2.CUDA, 0))96            else:97                devices = [blob_device_types[b] for b in ssa_i[0] + ssa_i[1]]98                assert all(d == devices[0] for d in devices)99                if devices[0] == "cuda":100                    op.device_option.CopyFrom(core.DeviceOption(caffe2_pb2.CUDA, 0))101 102    # update ops in predict_net103    predict_net_input_device_types = {104        (name, 0): _get_device_type(tensor)105        for name, tensor in zip(predict_net.external_input, tensor_inputs)106    }107    predict_net_device_types = infer_device_type(108        predict_net, known_status=predict_net_input_device_types, device_name_style="pytorch"109    )110    predict_net_ssa, _ = core.get_ssa(predict_net)111    _assign_op_device_option(predict_net, predict_net_ssa, predict_net_device_types)112 113    # update ops in init_net114    init_net_ssa, versions = core.get_ssa(init_net)115    init_net_output_device_types = {116        (name, versions[name]): predict_net_device_types[(name, 0)]117        for name in init_net.external_output118    }119    init_net_device_types = infer_device_type(120        init_net, known_status=init_net_output_device_types, device_name_style="pytorch"121    )122    _assign_op_device_option(init_net, init_net_ssa, init_net_device_types)123 124 125def export_caffe2_detection_model(model: torch.nn.Module, tensor_inputs: List[torch.Tensor]):126    """127    Export a caffe2-compatible Detectron2 model to caffe2 format via ONNX.128 129    Arg:130        model: a caffe2-compatible version of detectron2 model, defined in caffe2_modeling.py131        tensor_inputs: a list of tensors that caffe2 model takes as input.132    """133    model = copy.deepcopy(model)134    assert isinstance(model, torch.nn.Module)135    assert hasattr(model, "encode_additional_info")136 137    # Export via ONNX138    logger.info(139        "Exporting a {} model via ONNX ...".format(type(model).__name__)140        + " Some warnings from ONNX are expected and are usually not to worry about."141    )142    onnx_model = export_onnx_model(model, (tensor_inputs,))143    # Convert ONNX model to Caffe2 protobuf144    init_net, predict_net = Caffe2Backend.onnx_graph_to_caffe2_net(onnx_model)145    ops_table = [[op.type, op.input, op.output] for op in predict_net.op]146    table = tabulate(ops_table, headers=["type", "input", "output"], tablefmt="pipe")147    logger.info(148        "ONNX export Done. Exported predict_net (before optimizations):\n" + colored(table, "cyan")149    )150 151    # Apply protobuf optimization152    fuse_alias_placeholder(predict_net, init_net)153    if any(t.device.type != "cpu" for t in tensor_inputs):154        fuse_copy_between_cpu_and_gpu(predict_net)155        remove_dead_end_ops(init_net)156        _assign_device_option(predict_net, init_net, tensor_inputs)157    params, device_options = get_params_from_init_net(init_net)158    predict_net, params = remove_reshape_for_fc(predict_net, params)159    init_net = construct_init_net_from_params(params, device_options)160    group_norm_replace_aten_with_caffe2(predict_net)161 162    # Record necessary information for running the pb model in Detectron2 system.163    model.encode_additional_info(predict_net, init_net)164 165    logger.info("Operators used in predict_net: \n{}".format(_op_stats(predict_net)))166    logger.info("Operators used in init_net: \n{}".format(_op_stats(init_net)))167 168    return predict_net, init_net169 170 171def run_and_save_graph(predict_net, init_net, tensor_inputs, graph_save_path):172    """173    Run the caffe2 model on given inputs, recording the shape and draw the graph.174 175    predict_net/init_net: caffe2 model.176    tensor_inputs: a list of tensors that caffe2 model takes as input.177    graph_save_path: path for saving graph of exported model.178    """179 180    logger.info("Saving graph of ONNX exported model to {} ...".format(graph_save_path))181    save_graph(predict_net, graph_save_path, op_only=False)182 183    # Run the exported Caffe2 net184    logger.info("Running ONNX exported model ...")185    with ScopedWS("__ws_tmp__", True) as ws:186        ws.RunNetOnce(init_net)187        initialized_blobs = set(ws.Blobs())188        uninitialized = [inp for inp in predict_net.external_input if inp not in initialized_blobs]189        for name, blob in zip(uninitialized, tensor_inputs):190            ws.FeedBlob(name, blob)191 192        try:193            ws.RunNetOnce(predict_net)194        except RuntimeError as e:195            logger.warning("Encountered RuntimeError: \n{}".format(str(e)))196 197        ws_blobs = {b: ws.FetchBlob(b) for b in ws.Blobs()}198        blob_sizes = {b: ws_blobs[b].shape for b in ws_blobs if isinstance(ws_blobs[b], np.ndarray)}199 200        logger.info("Saving graph with blob shapes to {} ...".format(graph_save_path))201        save_graph(predict_net, graph_save_path, op_only=False, blob_sizes=blob_sizes)202 203        return ws_blobs204