Mismatch between Keras.predict and compiled TVM for imported Relay model

[chungs31]

Hello,

So I have been able to train binary AlexNet (1A1W) to top-1 accuracy to 43.2%, top-5 accuracy to 66.0%, fairly close to the published result. But I have been having issue getting the compiled model to have the same (or at least similar) output to Keras model.predict.

My setup:
Tensorflow version: 2.1.0
Keras version: 2.3.1
Target/Host: llvm/llvm (x86)

import tensorflow as tf
import numpy as np
import tvm
import tvm.relay as relay
import tvm.contrib.graph_runtime as runtime
from riptide.get_models import get_model
from riptide.binary.binary_layers import Config, DQuantize, XQuantize
from keras.applications.imagenet_utils import preprocess_input, decode_predictions
from keras.preprocessing import image

config = Config(actQ=DQuantize, weightQ=XQuantize, bits=1, use_act=True, use_bn=False, use_maxpool=True)
with config:
    model = get_model('alexnet')
test_input = tf.keras.Input(shape=[224, 224, 3], batch_size=1, dtype='float32')
output = model(test_input)
model.load_weights('...../alexnet_1A1W/model.ckpt-1003838')

# import from keras
mod, params = relay.frontend.from_keras(
  model,
  shape={'input_1': [1, 3, 224, 224]},
  layout='NCHW')

target = 'llvm'
target_host = 'llvm'
with relay.transform.build_config(opt_level=2):
  graph, lib, params = relay.build(mod, target=target, target_host=target_host, params=params)

ctx = tvm.cpu()

img = image.load_img('cat.jpeg', target_size=(224,224))
x = image.img_to_array(img, data_format="channels_first")
x = np.expand_dims(x, axis=0)
x = preprocess_input(x)

x_nhwc = image.img_to_array(img)
x_nhwc = np.expand_dims(x_nhwc, axis=0)
x_nhwc = preprocess_input(x_nhwc)
resnet = tf.keras.applications.ResNet50()

module = runtime.create(graph, lib, ctx)
module.set_input(**params)
module.set_input(model.input_names[0], x)
module.run()

modout = module.get_output(0)
print(modout.shape)

print("Keras prediction: ", decode_predictions(model.predict(x_nhwc)))
print("Keras ResNet50 prediction: ", decode_predictions(resnet.predict(x_nhwc)))
print("TVM prediction: ", decode_predictions(modout.asnumpy()))

This results in the following:

Keras prediction:  [[('n01924916', 'flatworm', 0.48209077), ('n02877765', 'bottlecap', 0.26524794), ('n01734418', 'king_snake', 0.051874492), ('n04209239', 'shower_curtain', 0.032657914), ('n09256479', 'coral_reef', 0.014033702)]]
Keras ResNet50 prediction:  [[('n02123045', 'tabby', 0.7215433), ('n02124075', 'Egyptian_cat', 0.21909083), ('n02123159', 'tiger_cat', 0.047373313), ('n03223299', 'doormat', 0.002071013), ('n02127052', 'lynx', 0.0011057281)]]
TVM prediction:  [[('n01980166', 'fiddler_crab', 0.5203081), ('n01914609', 'sea_anemone', 0.23222995), ('n11939491', 'daisy', 0.09443407), ('n02317335', 'starfish', 0.03952333), ('n02319095', 'sea_urchin', 0.017155001)]]

I should note that if I change the input image, model.predict will likely change but the TVM compiled prediction seems to always hover around fiddler_crab and sea_anemone. Am I doing something wrong here? I could expect that the model could predict incorrectly, but Keras and TVM to be "equally incorrect".