Export computer-vision model to ONNX

[saad-palapa]

Is your feature request related to a problem? Please describe. I would like to be able to export computer-vision models to ONNX

Describe the use case ONNX is a format that would be an addition to torchscript. It runs in many environments including, iOS, Android, web, and many more.

Describe the solution you'd like I wrote most of the code. I just need to test it and create a PR for it:

class LudwigWrapper(torch.nn.Module):
    def __init__(self, model):
        super(LudwigWrapper, self).__init__()
        self.model = model

    def forward(self, x):
        return self.model({"image_path": x})

def _export_classifier_onnx(model_path, export_path):
    ludwig_model = LudwigModel.load(model_path)
    model = LudwigWrapper(ludwig_model.model)  # Wrap the model
    model.eval()  # inference mode, is this needed.. I think onnx export does this for us

    width = ludwig_model.config["input_features"][0]["preprocessing"]["width"]
    height = ludwig_model.config["input_features"][0]["preprocessing"]["height"]
    example_input = torch.randn(1, 3, width, height, requires_grad=True)

    torch.onnx.export(
        model,
        example_input,
        export_path,
        opset_version=18,
        export_params=True,
        do_constant_folding=True,
        input_names=["input"],
        output_names=["combiner_hidden_1", "output", "combiner_hidden_2"],
    )

def _quantize(path_fp32, path_int8):
    from onnxruntime.quantization import quantize_dynamic

    quantize_dynamic(path_fp32, path_int8)  # type: ignore

Describe alternatives you've considered The alternative is to use other formats like CoreML

Additional context Join our slack channel: #computer-vision

[JaynouOliver]

If there is any issue, I would like to study about it and get started working on it, looks interesting to me, have previously worked on ONNX. can you provide some guidance and context?

[saad-palapa]

@JaynouOliver, we have a PR in the process: https://github.com/ludwig-ai/ludwig/pull/3761

If you would like to get involved in our computer vision efforts, join our #computer-vision slack channel.

[JaynouOliver]

Sure, already joined