# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/export_onnx.ipynb (unless otherwise specified).

__all__ = ['torch_to_onnx']

# Cell
import torch.nn as nn
import torch
from typing import Union
from pathlib import Path
from fastai.vision.all import *

import os
import onnx
import onnx.utils
from onnx import optimizer

# Cell
def torch_to_onnx(model:nn.Module,
                  activation:nn.Module=None,
                  save_path:str     = '../exported-models/',
                  model_fname:str   = 'onnx-model',
                  input_shape:tuple = (1,3,224,224),
                  input_name:str    = 'input_image',
                  output_names:Union[str,list] = 'output',
                  verbose:bool = True,
                  **export_args) -> os.PathLike:
    """
    Export a `nn.Module` -> ONNX
    This function exports the model with support for batching,
    checks that the export was done properly, and polishes the
    model up (removes unnecessary fluff added during conversion)

    The path to the saved model is returned
    Key Arguments
    =============
    * activation:  If not None, append this to the end of your model.
                   Typically a `nn.Softmax(-1)` or `nn.Sigmoid()`
    * input_shape: Shape of the inputs to the model
    """
    save_path = Path(save_path)
    if isinstance(output_names, str): output_names = [output_names]
    if activation: model = nn.Sequential(*[model, activation])
    model.eval()
    x = torch.randn(input_shape, requires_grad=True)
    x = x.cuda() if torch.cuda.is_available() else x
    model(x)
    dynamic_batch = {0: 'batch'}
    dynamic_axes  = {input_name : dynamic_batch}
    for out in output_names: dynamic_axes[out] = dynamic_batch
    torch.onnx._export(model, x, f"{save_path/model_fname}.onnx",
                       export_params=True, verbose=False,
                       input_names=[input_name], output_names=output_names,
                       dynamic_axes=dynamic_axes, keep_initializers_as_inputs=True,
                       **export_args)
    if verbose:
        print(f"Loading, polishing, and optimising exported model from {save_path/model_fname}.onnx")
    onnx_model = onnx.load(f'{save_path/model_fname}.onnx')
    model = onnx.utils.polish_model(onnx_model)
    #onnx.checker.check_model(model)

    # removing unused parts of the model
    passes = ["extract_constant_to_initializer", "eliminate_unused_initializer"]
    optimized_model = optimizer.optimize(onnx_model, passes)

    onnx.save(optimized_model, f'{save_path/model_fname}.onnx')
    print('<Exported ONNX model successfully>') # print regardless
    return f'{save_path/model_fname}.onnx'