search

microsoft / onnxscript

ONNX Script enables developers to naturally author ONNX functions and models using a subset of Python.
https://onnxscript.ai/
MIT License
285 stars 54 forks source link

[torchlib] Fix reflection_pad2d #964

Open justinchuby opened 1 year ago

justinchuby commented 1 year ago

Summary

ONNX Runtime raises [ONNXRuntimeError] : 1 : FAIL : Non-zero status code returned while running Pad node. Name:'_0x57e2840_n19' Status Message: Cannot use 'reflect' mode to pad dimension with a value of 0. Input shape:{0,3,3} when executing test ops_test.TestOutputConsistencyFullGraphCPU.test_output_match_opinfo__nn_functional_reflection_pad2d_constant_cpu_int64 in ONNX Script TorchLib.

To recreate this report, use

CREATE_REPRODUCTION_REPORT=1 python -m pytest onnxscript/tests/function_libs/torch_lib/ops_test.py -k test_output_match_opinfo__nn_functional_reflection_pad2d_constant_cpu_int64

To reproduce

import google.protobuf.text_format
import numpy as np
from numpy import array, float16, float32, float64, int32, int64
import onnx
import onnxruntime as ort

# Run n times
N = 1

onnx_model_text = """
ir_version: 8
producer_name: "pytorch"
producer_version: "2.1.0"
graph {
  node {
    input: "input_0"
    input: "input_1"
    output: "_val_2"
    name: "aten_reflection_pad2d_0"
    op_type: "aten_reflection_pad2d"
    doc_string: ""
    domain: "pkg.onnxscript.torch_lib"
  }
  name: "torch_jit"
  input {
    name: "input_0"
    type {
      tensor_type {
        elem_type: 7
        shape {
          dim {
            dim_value: 0
          }
          dim {
            dim_value: 3
          }
          dim {
            dim_value: 3
          }
        }
      }
    }
  }
  input {
    name: "input_1"
    type {
      tensor_type {
        elem_type: 7
        shape {
          dim {
            dim_value: 6
          }
        }
      }
    }
  }
  output {
    name: "_val_2"
    type {
      tensor_type {
        elem_type: 7
        shape {
          dim {
            dim_value: 2
          }
          dim {
            dim_value: 5
          }
          dim {
            dim_value: 5
          }
        }
      }
    }
  }
}
opset_import {
  domain: "pkg.onnxscript.torch_lib"
  version: 1
}
opset_import {
  domain: ""
  version: 18
}
functions {
  name: "aten_reflection_pad2d"
  input: "self"
  input: "padding"
  output: "return_val"
  node {
    output: "neg_1"
    name: "n0"
    op_type: "Constant"
    attribute {
      name: "value_ints"
      ints: -1
      type: INTS
    }
    domain: ""
  }
  node {
    output: "zero"
    name: "n1"
    op_type: "Constant"
    attribute {
      name: "value_ints"
      ints: 0
      type: INTS
    }
    domain: ""
  }
  node {
    input: "self"
    output: "tmp"
    name: "n2"
    op_type: "Shape"
    domain: ""
  }
  node {
    input: "tmp"
    output: "rank"
    name: "n3"
    op_type: "Size"
    domain: ""
  }
  node {
    output: "int64_2"
    name: "n4"
    op_type: "Constant"
    attribute {
      name: "value"
      t {
        data_type: 7
        int64_data: 2
        name: "int64_2"
      }
      type: TENSOR
    }
    domain: ""
  }
  node {
    input: "int64_2"
    input: "rank"
    output: "int64_2_cast"
    name: "n5"
    op_type: "CastLike"
    domain: ""
  }
  node {
    input: "rank"
    input: "int64_2_cast"
    output: "tmp_0"
    name: "n6"
    op_type: "Mul"
    domain: ""
  }
  node {
    input: "padding"
    output: "tmp_1"
    name: "n7"
    op_type: "Size"
    domain: ""
  }
  node {
    input: "tmp_0"
    input: "tmp_1"
    output: "tmp_2"
    name: "n8"
    op_type: "Sub"
    domain: ""
  }
  node {
    input: "tmp_2"
    input: "neg_1"
    output: "zero_count"
    name: "n9"
    op_type: "Reshape"
    domain: ""
  }
  node {
    input: "zero"
    input: "zero_count"
    output: "zeros"
    name: "n10"
    op_type: "Expand"
    domain: ""
  }
  node {
    input: "padding"
    input: "zeros"
    output: "torch_paddings"
    name: "n11"
    op_type: "Concat"
    attribute {
      name: "axis"
      i: 0
      type: INT
    }
    domain: ""
  }
  node {
    input: "torch_paddings"
    output: "size_d"
    name: "n12"
    op_type: "Size"
    domain: ""
  }
  node {
    output: "steps"
    name: "n13"
    op_type: "Constant"
    attribute {
      name: "value_ints"
      ints: -2
      type: INTS
    }
    domain: ""
  }
  node {
    input: "steps"
    input: "size_d"
    output: "ends"
    name: "n14"
    op_type: "Sub"
    domain: ""
  }
  node {
    input: "torch_paddings"
    input: "steps"
    input: "ends"
    input: "zero"
    input: "steps"
    output: "odd_elements"
    name: "n15"
    op_type: "Slice"
    domain: ""
  }
  node {
    input: "neg_1"
    input: "size_d"
    output: "ends_3"
    name: "n16"
    op_type: "Sub"
    domain: ""
  }
  node {
    input: "torch_paddings"
    input: "neg_1"
    input: "ends_3"
    input: "zero"
    input: "steps"
    output: "even_elements"
    name: "n17"
    op_type: "Slice"
    domain: ""
  }
  node {
    input: "odd_elements"
    input: "even_elements"
    output: "onnx_padding"
    name: "n18"
    op_type: "Concat"
    attribute {
      name: "axis"
      i: 0
      type: INT
    }
    domain: ""
  }
  node {
    input: "self"
    input: "onnx_padding"
    output: "return_val"
    name: "n19"
    op_type: "Pad"
    attribute {
      name: "mode"
      s: "reflect"
      type: STRING
    }
    domain: ""
  }
  doc_string: "reflection_pad2d(Tensor self, SymInt[4] padding) -> Tensor"
  opset_import {
    domain: ""
    version: 18
  }
  domain: "pkg.onnxscript.torch_lib"
}

"""

ort_inputs = {'input_0': array([], shape=(0, 3, 3), dtype=int64), 'input_1': array([1, 1, 1, 1, 1, 1])}

# Set up the inference session
session_options = ort.SessionOptions()
session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL
onnx_model = onnx.ModelProto()
google.protobuf.text_format.Parse(onnx_model_text, onnx_model)

# Uncomment this line to save the model to a file for examination
# onnx.save_model(onnx_model, "test_output_match_opinfo__nn_functional_reflection_pad2d_constant_cpu_int64.onnx")

onnx.checker.check_model(onnx_model)
session = ort.InferenceSession(onnx_model.SerializeToString(), session_options, providers=("CPUExecutionProvider",))

# Run the model
for _ in range(N):
    ort_outputs = session.run(None, ort_inputs)

Full error stack

[ONNXRuntimeError] : 1 : FAIL : Non-zero status code returned while running Pad node. Name:'_0x57e2840_n19' Status Message: Cannot use 'reflect' mode to pad dimension with a value of 0. Input shape:{0,3,3}
  File "/home/justinchu/dev/onnx-script/onnxscript/tests/function_libs/torch_lib/ops_test_common.py", line 533, in _capture_graph_and_evaluate_torch_script_evaluator
    return _safe_ort_session_run(onnx_model.SerializeToString(), ort_inputs)
  File "/home/justinchu/dev/onnx-script/onnxscript/tests/function_libs/torch_lib/ops_test_common.py", line 349, in _safe_ort_session_run
    raise return_dict["error"]

Environment

OS: Linux-5.15.0-1042-azure-x86_64-with-glibc2.35
Python version: 3.10.9 (main, Jan 11 2023, 15:21:40) [GCC 11.2.0]
onnx==1.15.0.dev20230731
onnxruntime==1.15.1
numpy==1.25.1
torch==2.1.0.dev20230622+cpu
xiaowuhu commented 1 year ago

In OnnxScript, the simlar test case is:

ort_inputs = {'input_0': array([], shape=(0, 3, 3), dtype=int64), 'input_1': array([1,1])}

Inside the function, the [1,1] will be converted to [0, 0, 1, 0, 0, 1], so the op.Pad() can recognize.

But here the test case is:

ort_inputs = {'input_0': array([], shape=(0, 3, 3), dtype=int64), 'input_1': array([1, 1, 1, 1, 1, 1])}

In input_1, the first and fourth number should be 0, because input_0's first dimension size is 0. And, 6 numbers for padding value only apply to 4d tensor. No idea why the test case here is different.