Brecht De Cock
Brecht De Cock

Reputation: 1

Inferencing ONNX model of pytorch's mask r-cnn never works

The problem: inferencing onnx model gives empty results or weirdly shaped results

What i'm trying: pytorch pretrained mask-rcnn -> finetune on dataset -> save as onnx -> inference on onnx-> plot results

Everything i currently have works until the inference. my main.py file: https://pastebin.com/3jNfZdBi

getting info about my saved onnx model gives

Model Input Info:
Name: input
Shape: ['batch_size', 3, 'height', 'width']
Type: tensor(float)

Model Output Info:
Name: boxes
Shape: ['Concatboxes_dim_0', 4]
Type: tensor(float)
Name: labels
Shape: ['Gatherlabels_dim_0']
Type: tensor(int64)
Name: scores
Shape: ['Gatherlabels_dim_0']
Type: tensor(float)
Name: masks
Shape: ['Unsqueezemasks_dim_0', 'Unsqueezemasks_dim_1', 'Unsqueezemasks_dim_2', 'Unsqueezemasks_dim_3']
Type: tensor(float)

My code for inferencing: https://pastebin.com/wxvp649G

I suspect i either: save things wrongly to onnx OR don't preprocess my data correctly OR my inferencing code is wrong (or something other i dont know about)

Code for saving to onnx

def save_model_onnx(models_file_path, model, torch_input):
    # Traditional export method. There is also an experimental dynamo_export method
    torch.onnx.export(
        model.cpu(),
        torch_input.cpu(),
        models_file_path, #full path to the model including the model itself i.e. ./models/model.onnx
        export_params = True,
        opset_version=15,  # Choose a supported ONNX opset version
        do_constant_folding=True,  # Fold constant nodes for optimization
        input_names = ['input'],
        output_names = ['boxes', 'labels', 'scores', 'masks'],
        dynamic_axes={
        "input": {0: "batch_size", 2: "height", 3: "width"},  
        }
    )
    logging.info(f"Model saved at {models_file_path}")

Upvotes: 0

Views: 43

Answers (0)

Related Questions