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。
Here, I used GravNetConv blocks in the model, so I encountered some problems with custom operators. As follows
# 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)
GravNetConv(
in_channels=dim1,
out_channels=dim1 * 2,
space_dimensions=space_dimensions,
k=k,
propagate_dimensions=dim1,
),
This is the code of my onnx everywhere and the inference using onnxruntime
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
The following is the error report
===Start exporting ONNX ===
The model has been created.
The weights have been loaded.
dummy_x shape: torch.Size([50, 7]), dtype: torch.float32
dummy_batch shape: torch.Size([50]), dtype: torch.int64
Verify forward propagation…
Output shape: torch.Size([50, 11])
Verification passed!
Start exporting ONNX…
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.)
_C._jit_pass_onnx_remove_inplace_ops_for_onnx(graph, module)
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.
return fn(g, to_cast_func(g, input, False), to_cast_func(g, other, False))
[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)
[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)
[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)
[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)
Export successful!
The ONNX model verification was successful!
Verification failure: [ONNXRuntimeError] : 1: FAIL: Load model from cdcnet.onnx failed:Fatal error: CustomDomain:KNN(-1) is not a registered function/op
Traceback (most recent call last):
File “c:\ Code \ Dataset \onnx_model\onnxtest.py”, line 289, in
validate_onnx(“cdcnet.onnx”)
File “c:\ Code \ Dataset \onnx_model\onnxtest.py”, line 265, in validate_onnx
ort_session = ort.InferenceSession(onnx_path)
File “D:\anaconda\envs\onnx\lib\site-packages\onnxruntime\capi\onnxruntime_inference_collection.py”, line 419, in init
self._create_inference_session(providers, provider_options, disabled_optimizers)
File “D:\anaconda\envs\onnx\lib\site-packages\onnxruntime\capi\onnxruntime_inference_collection.py”, line 480, in _create_inference_session
sess = C.InferenceSession(session_options, self._model_path, True, self._read_config_from_model)
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