ONNX export failed: Couldn't export Python operator CorrSampler

[sunmooncode]

When I converted the ONNX model, I encountered that the CorrSampler could not be converted. Can you provide some suggestions?

Looking forward to your reply！

[lahavlipson]

CorrSampler is an optional faster implementation of the correlation sampler, enabled using --corr_implementation reg_cuda. If you use --corr_implementation reg, do you still get the same issue?

Sorry if my question doesn't make sense, I am not so familiar with ONNX.

[sunmooncode]

thank you for your reply！
when I use --corr_implementation reg, Here is another problem！

RuntimeError: Exporting the operator grid_sampler to ONNX opset version 12 is not supported. Please open a bug to request ONNX export support for the missing operator.

[lahavlipson]

It looks like ONNX is not yet compatible with torch.nn.functional.grid_sample, however someone created an alternate implementation to address this issue:

pytorch/pytorch#27212 (comment)

[daigo0927]

This is correct (strictly, the ONNX opset-v16 offers this operation but PyTorch doesn't reflect it yet).

It looks like ONNX is not yet compatible with torch.nn.functional.grid_sample

I recommend to use contrib-op provided by Microsoft. This enables you to export a model with F.grid_sample opeation to ONNX.

from torch.onnx import register_custom_op_symbolic
import torch.onnx.symbolic_helper as sym_help


def grid_sampler(g, input, grid, mode, padding_mode, align_corners):
    # mode
    #   'bilinear'      : onnx::Constant[value={0}]
    #   'nearest'       : onnx::Constant[value={1}]
    #   'bicubic'       : onnx::Constant[value={2}]
    # padding_mode
    #   'zeros'         : onnx::Constant[value={0}]
    #   'border'        : onnx::Constant[value={1}]
    #   'reflection'    : onnx::Constant[value={2}]
    mode = sym_help._maybe_get_const(mode, "i")
    padding_mode = sym_help._maybe_get_const(padding_mode, "i")
    mode_str = ['bilinear', 'nearest', 'bicubic'][mode]
    padding_mode_str = ['zeros', 'border', 'reflection'][padding_mode]
    align_corners = int(sym_help._maybe_get_const(align_corners, "b"))

    return g.op("com.microsoft::GridSample", input, grid,
                mode_s=mode_str,
                padding_mode_s=padding_mode_str,
                align_corners_i=align_corners)


# Register contrib GridSample operation as a counterpart for F.grid_sample
register_custom_op_symbolic('::grid_sampler', grid_sampler, 1)

torch.onnx.export(
    ... # a model with F.grid_sample 
)

See my gist for a simple usage ;)

[daviszhu88]

When I converted the ONNX model, I encountered that the CorrSampler could not be converted. Can you provide some suggestions?

Looking forward to your reply！

Hi, did you convert the pth model to onnx format? Could you kindly share your script if you do? Thanks.