Wrong output shape from a graph generated by nnvm.build.compiler

I am trying to get a better understanding of the TVM software stack with the below experiment Tensorflow code.

The code below simply adds the left half, with the right half of the image. However, after using nnvm.compiler.build to build the graph, the graph returns output of shape (1920, 3) instead of (1080, 1920, 3). Did I miss any important step?

# import statements ...
img = cv2.imread('image.jpg', cv2.COLOR_BGR2RGB)
height, width, channel = img.shape  # (1080, 1920, 3)

img_1 = np.copy(img)  # For x, right half image
img_1[:, :(width // 2), :] = 0

img_2 = np.copy(img)  # For b, left half image
img_2[:, (width // 2):, :] = 0

# graph: y = x + b
x = tf.placeholder(dtype=tf.float32, shape=(1080, 1920, 3), name='x')
b = tf.constant(img_2, dtype=tf.float32, shape=(1080, 1920, 3), name='b')
y = tf.math.add(x, b, name='output')

# Save graph to 'path_to_file.pb' ...

with tf.gfile.GFile('path_to_file.pb', 'rb') as f:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())

# target, target_host settings ...
with tf.Graph().as_default() as graph:
    tf.import_graph_def(graph_def, name='')
    sym, params = nnvm.frontend.from_tensorflow(graph_def)

    shape_dict = {'x': (1080, 1920, 3)}
    dtype_dict = {'x': 'float32'}

    graph, lib, params = nnvm.compiler.build(graph = sym,
                                             shape = shape_dict,
                                             dtype = dtype_dict,
                                             target = target,
                                             params = params,
                                             target_host=target_host)

    m = graph_runtime.create(graph, lib, ctx)
    m.set_input('x', tvm.nd.array(img_1))
    m.set_input(**params)  # <--- Not sure what this line do.
    m.run()

    # The below line returns output of shape (1920, 3) <--- ?
    output = m.get_output(0).asnumpy()[0]

I am not familiar with the tensorflow importer, but I am curious if some shapes may truncated because the usual case is a NCHW data layout, and HWC is not properly supported. It is not really a proper fix, but can you try importing the images as (1, 3, 1080, 1920)?

Layout shouldn’t matter as this is just an add operation. let me check

@eqy Yes, I tried to convert all values to 4D tensor and it works now.

2 Likes