export.py

"""Export a YOLOv5 *.pt model to ONNX formats
     Modified version of export.py from yolov5 repository/ commit 63dd65e7edd96debbefa81e22f3d5cfb07dd2ba4 

   ADDED: 
      - Function export_tf() to convert ONNX to TF representation using onnx_tf
      - Function export_tflite() to convert TF representation to TFLite using TF converter
      - Function export_openvino() to convert ONNX to OpenVINO IR representaion in docker (Docker must be installed!)

Usage:
    $ python export_models.py --weights models/yolov5s.pt --img 640 --batch 1
"""

import argparse
import sys
import time
from pathlib import Path

import torch
import torch.nn as nn
from torch.utils.mobile_optimizer import optimize_for_mobile

sys.path.insert(0, './yolov5')

FILE = Path(__file__).absolute()
sys.path.append(FILE.parents[0].as_posix())  # add yolov5/ to path

from models.common import Conv
from models.yolo import Detect
from models.experimental import attempt_load
from utils.activations import Hardswish, SiLU
from utils.general import colorstr, check_img_size, set_logging, check_requirements#, file_size -> no such function in v5.0
from utils.torch_utils import select_device


def file_size(file):
    '''
    Function from utils.general import file size 
    It is not present in the stable release v5.0
    '''
    # Return file size in MB
    return Path(file).stat().st_size / 1e6


def export_onnx(model, img, file, opset, train, dynamic, simplify):
    # ONNX model export
    prefix = colorstr('ONNX:')
    try:
        check_requirements(('onnx', 'onnx-simplifier'))
        import onnx

        print(f'\n{prefix} starting export with onnx {onnx.__version__}...')
        f = file.with_suffix('.onnx')
        torch.onnx.export(model, img, f, verbose=False, opset_version=opset,
                          training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL,
                          do_constant_folding=not train,
                          input_names=['images'],
                          output_names=['output'],
                          dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'},  # shape(1,3,640,640)
                                        'output': {0: 'batch', 1: 'anchors'}  # shape(1,25200,85)
                                        } if dynamic else None)

        # Checks
        model_onnx = onnx.load(f)                               # load onnx model
        onnx.checker.check_model(model_onnx)                    # check onnx model

        # Simplify
        if simplify:
            try:
                import onnxsim

                print(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')
                model_onnx, check = onnxsim.simplify(
                    model_onnx,
                    dynamic_input_shape=dynamic,
                    input_shapes={'images': list(img.shape)} if dynamic else None)
                assert check, 'assert check failed'
                onnx.save(model_onnx, f)
            except Exception as e:
                print(f'{prefix} simplifier failure: {e}')
        print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')

    except Exception as e:
        print(f'{prefix} export failure: {e}')


def export_tf(file):
    '''
    Convert model from ONNX to TensoFlow representation (.pb) using onnx_tf
    '''
    import onnx
    import onnx_tf

    f = file.with_suffix('.onnx')

    model_onnx = onnx.load(f)

    prefix = colorstr('ONNX_TF:')
    print(f'\n{prefix} starting export with onnx_tf {onnx_tf.__version__}...')
    try:
        tf_rep = onnx_tf.backend.prepare(model_onnx, strict=False)   # prepare tf representation
        tf_rep.export_graph('models')                                # export the model graph
        print(f'{prefix} export success, saved as {f}')
    except Exception as e:
        print(f'{prefix} export failure: {e}')


def export_tflite(file):
    '''
    Convert model from TensoFlow representation (.pb) to TensorFlow Lite (.tflite)
    '''
    import tensorflow as tf

    f = file.with_suffix('.tflite')

    prefix = colorstr('TensorFlow:')
    print(f'\n{prefix} starting export with TensorFlow {tf.__version__}...')
    try:
        # Convert the model
        converter = tf.lite.TFLiteConverter.from_saved_model('models')
        # TFLite builtin operator library only supports a limited number of TF operators 
        converter.target_spec.supported_ops = [
          tf.lite.OpsSet.TFLITE_BUILTINS,        # enable TensorFlow Lite ops.
          tf.lite.OpsSet.SELECT_TF_OPS           # enable TensorFlow ops.
        ]
        tflite_model = converter.convert()
        
        open(f, "wb").write(tflite_model)
    except Exception as e:
        print(f'{prefix} export failure: {e}')

def export_openvino(file):
    '''
    Convert model from ONNX to OpenVINO representation using docker
    '''
    import subprocess

    f = file.with_suffix('.onnx')

    prefix = colorstr('ONNX_VINO:')

    docker_command = 'docker run --rm -v $PWD/models:/home/openvino/models --user "$(id -u):$(id -g)" -w /home/openvino openvino/ubuntu20_dev:latest ' + '/bin/bash -c "python3 /opt/intel/openvino/deployment_tools/model_optimizer/mo.py --progress --input_shape [1,3,320,320] --input_model models/yolov5s.onnx --output_dir models/yolov5_openvino --data_type half"'

    print(f'\n{prefix} starting export with openvino...')
    try:
        with open("docker.log", "a") as output:
           subprocess.call(docker_command, shell=True, stdout=output, stderr=output)
        print(f'{prefix} export success, saved as {f}')
    except Exception as e:
        print(f'{prefix} export failure: {e}')


def run(weights='./yolov5s.pt',                     # weights path
        img_size=(640, 640),                        # image (height, width)
        batch_size=1,                               # batch size
        device='cpu',                               # cuda device, i.e. 0 or 0,1,2,3 or cpu
        include=('onnx','tf','tflite','openvino'),  # include formats
        half=False,                                 # FP16 half-precision export
        inplace=False,                              # set YOLOv5 Detect() inplace=True
        train=False,                                # model.train() mode
        optimize=False,                             # TorchScript: optimize for mobile
        dynamic=False,                              # ONNX: dynamic axes
        simplify=False,                             # ONNX: simplify model
        opset=12,                                   # ONNX: opset version
        ):
    t = time.time()
    include = [x.lower() for x in include]

    img_size *= 2 if len(img_size) == 1 else 1  # expand
    file = Path(weights)

    # Load PyTorch model
    device = select_device(device)
    assert not (device.type == 'cpu' and half), '--half only compatible with GPU export, i.e. use --device 0'
    model = attempt_load(weights, map_location=device)  # load FP32 model
    names = model.names

    # Input
    gs = int(max(model.stride))  # grid size (max stride)
    img_size = [check_img_size(x, gs) for x in img_size]  # verify img_size are gs-multiples
    img = torch.zeros(batch_size, 3, *img_size).to(device)  # image size(1,3,320,192) iDetection

    # Update model
    if half:
        img, model = img.half(), model.half()  # to FP16
    model.train() if train else model.eval()   # training mode = no Detect() layer grid construction
    for k, m in model.named_modules():
        if isinstance(m, Conv):  # assign export-friendly activations
            if isinstance(m.act, nn.Hardswish):
                m.act = Hardswish()
            elif isinstance(m.act, nn.SiLU):
                m.act = SiLU()
        elif isinstance(m, Detect):
            m.inplace = inplace
            m.onnx_dynamic = dynamic
            # m.forward = m.forward_export  # assign forward (optional)

    for _ in range(2):
        y = model(img)  # dry runs
    print(f"\n{colorstr('PyTorch:')} starting from {weights} ({file_size(weights):.1f} MB)")

    ## TODO: need to be more clear. ONNX must be run first to convert to TensorFlow
    #if 'tf' in include and 'onnx' not in include:
    #    include.append('onnx')

    # Exports
    if 'onnx' in include:
        export_onnx(model, img, file, opset, train, dynamic, simplify)
    if 'tf' in include:
        export_tf(file)
    if 'tflite' in include:
        export_tflite(file)
    if 'openvino' in include:
        export_openvino(file)

    # Finish
    print(f'\nExport complete ({time.time() - t:.2f}s). Visualize with https://github.com/lutzroeder/netron.')


def parse_opt():
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', type=str, default='models/yolov5s.pt', help='weights path')
    parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image (height, width)')
    parser.add_argument('--batch-size', type=int, default=1, help='batch size')
    parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--include', nargs='+', default=['onnx', 'tf', 'tflite','openvino'], help='include formats')
    parser.add_argument('--half', action='store_true', help='FP16 half-precision export')
    parser.add_argument('--inplace', action='store_true', help='set YOLOv5 Detect() inplace=True')
    parser.add_argument('--train', action='store_true', help='model.train() mode')
    parser.add_argument('--optimize', action='store_true', help='TorchScript: optimize for mobile')
    parser.add_argument('--dynamic', action='store_true', help='ONNX: dynamic axes')
    parser.add_argument('--simplify', action='store_true', help='ONNX: simplify model')
    parser.add_argument('--opset', type=int, default=12, help='ONNX: opset version')
    opt = parser.parse_args()
    return opt


def main(opt):
    set_logging()
    print(colorstr('export: ') + ', '.join(f'{k}={v}' for k, v in vars(opt).items()))
    run(**vars(opt))


if __name__ == "__main__":
    opt = parse_opt()
    main(opt)