How to get the INPUT and OUTPUT of relay call node

Questions
1

Hi All,

How can I get (I guess I need to write a simple relay pass) INPUT(s), INPUT shapes and OUTPUT(s), OUTPUT shapes of a call node?

Let us say I have the following example, and I want to get the inputs and outputs? I guess the input is simple, but I was not able to figure out how to get the output (data type and shape)?

`

Set the dimensions

dim1 = 20
dim2 = 1
#
dshape = (dim1, dim2)
x = relay.var('x', shape=dshape)
y = relay.var('y', shape=dshape)

x2 = relay.var('x2', shape=dshape)
y2 = relay.var('y2', shape=dshape)

# f = relay.Function([x, y], x + y)
add = relay.op.add(x, y)
mul = relay.op.multiply(add, add)
result1 = relay.sum(mul, 0)

`

This is what result1 looks like.
print(result1) free_var %x: Tensor[(20, 1), float32]; free_var %y: Tensor[(20, 1), float32]; %0 = add(%x, %y); %1 = multiply(%0, %0); sum(%1, axis=[0])

2

You need to infer type to get the output type. The following script is based on yours:

import tvm
from tvm import relay
from tvm.relay import transform

dim1 = 20
dim2 = 1
#
dshape = (dim1, dim2)
x = relay.var('x', shape=dshape)
y = relay.var('y', shape=dshape)

x2 = relay.var('x2', shape=dshape)
y2 = relay.var('y2', shape=dshape)

add = relay.op.add(x, y)
mul = relay.op.multiply(add, add)
result1 = relay.sum(mul, 0)
func = relay.Function(relay.analysis.free_vars(result1), result1)
mod = tvm.IRModule.from_expr(func)

try:
    print(mod["main"].checked_type)
except ValueError as err:
    print(err)
    # The type checker has not populated the checked_type for this node
    pass

mod = transform.InferType()(mod)
print(mod["main"].checked_type)
# FuncType([], [TensorType([20, 1], float32), TensorType([20, 1], float32)], TensorType([1], float32), [])
1 Like
3

Thanks, @comaniac .

I have written the following script to show what I mean. I can get the “INPUT(s)” by using mod[“main”].checked_type.arg_types, but it returns a typle that includes my INPUT as well as additional tuples as shown below. Assume, that I have an IRModule which I do not know how many inputs are there, then how can I know which ones are my inputs?

def get_inputs_to_IRModule(mod):
    try:
        print(mod["main"].checked_type)
    except ValueError as err:
        print(err)
        # The type checker has not populated the checked_type for this node
        pass

    mod = transform.InferType()(mod)
    # print(mod["main"].checked_type)
    # print(mod["main"].checked_type.arg_types)

    return mod["main"].checked_type.arg_types


if True:
    onnxfile="mnist/model.onnx"
    # load ONNX model
    onnxmodel =  onnx.load_model(onnxfile)

    input_tensor = numpy_helper.to_array(onnx.load_tensor("mnist/test_data_set_0/input_0.pb"))

    input_name = onnxmodel.graph.input[0].name
    shape_dict = {input_name: input_tensor.shape}
    print("------Inputs-----: {}".format(shape_dict))

    # convert ONNX model to IR module. THIS LINE gives an error.
    #dtype = 'float32'
    mod, params = relay.frontend.from_onnx(onnxmodel, shape_dict)
    print("-----Inputs to MNIST are (Gives more than what I want)-----:{}". format(get_inputs_to_IRModule(mod)))
4

The short answer is no. From an IRModule’s point of view, all inputs, including your data input and the weight parameter inputs, are the “inputs” of this module, and they are no difference. From your example, I would say shape_dict already has the information you’re looking for.

5

Thanks @comaniac – Yes, that is what I am getting (all inputs + weights).

The reason I am seeking (the code above does not represent well) is because I have a “RANDOM” IRModule that I do not know the INPUT info (size, shape). Basically, what I mean by this is I do not have access to shape_dict as above.

I guess there is a “NOT CLEAN” way to get the inputs, but if you have suggestions please let me know.

6

I was thinking about this a little bit. Let me present the problem definition one more time, and if it helps anyone has suggestions.

Assumptions: I have an IRModule in which I do not know the input shape and number of inputs. My problem is I need to figure out how many inputs are there, and what are the shapes of input.

I tried to get the inputs for my “unknown input IRModule” with various ways, and I also tried to get the inputs of equivalent relay expression for my IRModule.

Both attempts failed with giving me inputs + weights as you mentioned by Cody @comaniac. Any suggestions?

  1. IRModule ==> All inputs = inputs + weights
  2. Relay Expression (relay function) = => All inputs = inputs + weights