我在onnx文件导出后，出现这个问题：onnxruntime.capi.onnxruntime_pybind11_state.Fail: [ONNXRuntimeError] : 1 : FAIL : Load model from cdcnet.onnx failed:Fatal error: CustomDomain:KNN(-1) is not a registered function/op。

Mio: <p>I’m a novice in this research and hope to receive responses from all the teachers, who can provide me with a solution. Thank you very much for your sincere blessings。</p>
<p><strong>Here, I used GravNetConv blocks in the model, so I encountered some problems with custom operators. As follows</strong></p>
<pre><code class="lang-auto"># Custom operator symbol function
@parse_args("v", "v", "i", "v", "v", "b", "i")
def knn_symbolic(g, x, y, k, batch_x, batch_y, cosine, num_workers):
# Obtain the shape of the input tensor x
x_shape = g.op("Shape", x)

# Ensure that x is two-dimensional
x_rank = g.op("Size", x_shape)
x_rank = g.opp ("Cast", x_rank, to_i=6) # Convert to int64

# If x is not two-dimensional, add one dimension
if x_rank ! = 2:
axes = g.op("Constant", value_t=torch.tensor([0], dtype=torch.int64))
x = g.op("Unsqueeze", x, axes)

# Call the custom KNN operator
knn_node = g.op(
"CustomDomain::KNN",
x,
y,
batch_x,
batch_y,
k_i=k,
cosine_i=int(cosine),
num_workers_i=num_workers,
outputs=1
)

return knn_node

@parse_args("v", "v", "i", "i", "i")
def gravnet_conv_symbolic(g, x, batch, k, space_dimensions, propagate_dimensions):
# Obtain the shape of the input tensor x
x_shape = g.op("Shape", x)

# Ensure that x is two-dimensional
x_rank = g.op("Size", x_shape)
x_rank = g.opp ("Cast", x_rank, to_i=6) # Convert to int64

# If x is not two-dimensional, add one dimension
if x_rank ! = 2:
axes = g.op("Constant", value_t=torch.tensor([0], dtype=torch.int64))
x = g.op("Unsqueeze", x, axes)

# Call the custom GravNetConv operator
gravnet_node = g.op(
"CustomDomain::GravNetConv",
x,
batch,
k_i=k,
space_dimensions_i=space_dimensions,
propagate_dimensions_i=propagate_dimensions,
outputs=1
)

return gravnet_node

@parse_args("v", "i", "v", "v", "s", "b")
def custom_scatter_reduce(g, self, dim, index, src, reduce, include_self):
if not include_self:
scatter_node = g.op(
"CustomDomain::CustomScatterReduce",
self,
index,
src,
dim_i=dim,
reduce_s=reduce,
include_self_i=True,
outputs=1
)
else:
reduce_mode = {
"sum": "add",
"prod": "mul",
"amin": "min",
"amax": "max",
}
onnx_reduce = reduce_mode.get(reduce, "add")
scatter_node = g.op(
"ScatterElements",
self,
index,
src,
axis_i=dim,
reduction_s=onnx_reduce,
outputs=1
)

# Add shape reasoning
self_shape = g.op("Shape", self)
scatter_node = g.op("Cast", scatter_node, to_i=1) # Ensure type consistency
scatter_node = g.op("Reshape", scatter_node, self_shape)

return scatter_node

# Register custom operators
register_custom_op_symbolic("torch_cluster::knn", knn_symbolic, 14)
register_custom_op_symbolic("torch_geometric::gravnet_conv", gravnet_conv_symbolic, 14)
register_custom_op_symbolic("aten::scatter_reduce", custom_scatter_reduce, 14)
</code></pre>
<pre><code class="lang-auto">GravNetConv(
                                    in_channels=dim1,
                                    out_channels=dim1 * 2,
                                    space_dimensions=space_dimensions,
                                    k=k,
                                    propagate_dimensions=dim1,
                                ),
</code></pre>
<p><strong>This is the code of my onnx everywhere and the inference using onnxruntime</strong></p>
<pre><code class="lang-auto">def export_onnx(model_path, onnx_path, input_dim=7, batch_size=50):
try:
print("=== Start ONNX export ===")

# Model Initialization
model = CDCNet(input_dim=input_dim)
print(" Model Created ")

# Loading Weight
model = load_weights(model, model_path)
model.eval()
print(" Weights have been loaded ")

# Create the input tensor
dummy_x = torch.randn(batch_size, input_dim)
dummy_batch = torch.zeros(batch_size, dtype=torch.long)

print(f"dummy_x shape: {dummy_x.shape}, dtype: {dummy_x.dtype}")
print(f"dummy_batch shape: {dummy_batch.shape}, dtype: {dummy_batch.dtype}")

# Verify forward propagation
print(" Verify forward propagation..." )
with torch.no_grad():
output = model(dummy_x, dummy_batch)
print(f" output shape: {output.shape}")
assert output.shape[1] == 11 # Assume the output dimension is 11
print(" Verified!" )

# Export ONNX
print(" Start exporting ONNX..." )
torch.onnx.export(
model,
(dummy_x, dummy_batch),
onnx_path,
opset_version=14,
input_names=["x", "batch"],
output_names=["output"],
dynamic_axes={
"x": {0: "batch_size"},
"batch": {0: "batch_size"},
"output": {0: "batch_size"}
},
verbose=False
)
print(" Export successful!" )

except Exception as e:
print(f" Export failed: {str(e)}")
raise

def validate_onnx(onnx_path):
try:
# Load the ONNX model
onnx_model = onnx.load(onnx_path)
onnx.checker.check_model(onnx_model)
print("ONNX model verification successful!" )

# Inference using ONNX Runtime
ort_session = ort.InferenceSession(onnx_path)
dummy_x = np.random.randn(50, 7).astype(np.float32)
dummy_batch = np.zeros(50, dtype=np.int64)
ort_inputs = {
"x": dummy_x,
"batch": dummy_batch
}
ort_output = ort_session.run(None, ort_inputs)

# Compare the outputs of PyTorch and ONNX
with torch.no_grad():
model = CDCNet(input_dim=7)
model = load_weights(model, "best_model.pt")
model.eval()
torch_output = model(torch.tensor(dummy_x), torch.tensor(dummy_batch)).numpy()

print(" Whether the ONNX output is consistent with the PyTorch output :", np.allclose(torch_output, ort_output[0], atol=1e-4))

except Exception as e:
print(f" Verification failed: {str(e)}")
raise
</code></pre>
<p><strong>The following is the error report</strong></p>
<p>===Start exporting ONNX ===<br>
The model has been created.<br>
The weights have been loaded.<br>
dummy_x shape: torch.Size([50, 7]), dtype: torch.float32<br>
dummy_batch shape: torch.Size([50]), dtype: torch.int64<br>
Verify forward propagation…<br>
Output shape: torch.Size([50, 11])<br>
Verification passed!<br>
Start exporting ONNX…<br>
D:\anaconda\envs\onnx\lib\site-packages\torch\onnx\utils.py:609: UserWarning:  ONNX Preprocess - Removing mutation from node aten::scatter_add_ on block input:  ‘batch’. This changes graph semantics. (Triggered internally at  C:\actions-runner_work\pytorch\pytorch\pytorch\torch\csrc\jit\passes\onnx\remove_inplace_ops_for_onnx.cpp:353.)<br>
_C._jit_pass_onnx_remove_inplace_ops_for_onnx(graph, module)<br>
D:\anaconda\envs\onnx\lib\site-packages\torch\onnx\symbolic_opset9.py:1959: FutureWarning:  torch.onnx.symbolic_opset9._cast_Bool’ is deprecated in version 2.0 and will be removed in the future. Please Avoid using this function and create a Cast node instead.<br>
return fn(g, to_cast_func(g, input, False), to_cast_func(g, other, False))<br>
[W420 14:01:34.000000000 shape_type_inference.cpp:1999] Warning: The shape inference of CustomDomain::KNN type is missing,  so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function.  (function UpdateReliable)<br>
[W420 14:01:34.000000000 shape_type_inference.cpp:1999] Warning: The shape inference of CustomDomain::KNN type is missing,  so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function.  (function UpdateReliable)<br>
[W420 14:01:35.000000000 shape_type_inference.cpp:1999] Warning: The shape inference of CustomDomain::KNN type is missing,  so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function.  (function UpdateReliable)<br>
[W420 14:01:35.000000000 shape_type_inference.cpp:1999] Warning: The shape inference of CustomDomain::KNN type is missing,  so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function.  (function UpdateReliable)<br>
Export successful!<br>
The ONNX model verification was successful!<br>
Verification failure: [ONNXRuntimeError] : 1: FAIL: Load model from cdcnet.onnx failed:Fatal error: CustomDomain:KNN(-1) is not a registered function/op<br>
Traceback (most recent call last):<br>
File “c:\ Code \ Dataset \onnx_model\onnxtest.py”, line 289, in <br>
validate_onnx(“cdcnet.onnx”)<br>
File “c:\ Code \ Dataset \onnx_model\onnxtest.py”, line 265, in validate_onnx<br>
ort_session = ort.InferenceSession(onnx_path)<br>
File “D:\anaconda\envs\onnx\lib\site-packages\onnxruntime\capi\onnxruntime_inference_collection.py”, line 419,  in <strong>init</strong><br>
self._create_inference_session(providers, provider_options, disabled_optimizers)<br>
File “D:\anaconda\envs\onnx\lib\site-packages\onnxruntime\capi\onnxruntime_inference_collection.py”, line 480,  in _create_inference_session<br>
sess = C.InferenceSession(session_options, self._model_path, True, self._read_config_from_model)<br>
onnxruntime.capi.onnxruntime_pybind11_state.Fail: [ONNXRuntimeError] : 1 : FAIL :  Load model from cdcnet.onnx failed:Fatal error: CustomDomain:KNN(-1) is not a registered function/op</p>