jwfromm / riptide Goto Github PK

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()))

def _conv(n, co, h, w, vh, vw, vc):
        b1b2 = (b1+b2).astype(out_dtype)
        if unipolar:
            return tvm.sum((tvm.popcount(
                data_vec[n, h, w, ci, vh*HSTR+dh, vw*WSTR+dw, b1].astype(out_dtype) &
                kernel_vec[co, ci, dh, dw, b2, vc].astype(out_dtype))  -
                            tvm.popcount(
                                data_vec[n, h, w, ci, vh*HSTR+dh, vw*WSTR+dw, b1].astype(out_dtype)
                                & ~kernel_vec[co, ci, dh, dw, b2, vc]).astype(out_dtype)) << b1b2,
                           axis=[ci, dh, dw, b1, b2])

        return tvm.sum((tvm.popcount(
            data_vec[n, h, w, ci, vh*HSTR+dh, vw*WSTR+dw, b1] &
            kernel_vec[co, ci, dh, dw, b2, vc])).astype(out_dtype) << b1b2,
                       axis=[ci, dh, dw, b1, b2])

import tvm
import numpy as np
import topi
from tvm import relay
import topi.testing
from tvm.contrib import graph_runtime
from topi.util import get_const_tuple

batch = 1
in_height = in_width = in_size = 32
in_dim = 32
out_dim = 32
in_channel = 32
num_filter = 32
kernel = 3
stride = (1, 1)
padding = (1, 1)
activation_bits = 1
weight_bits = 1
unipolar = True

input_dtype = 'uint8'
out_dtype = 'int8'

def generate_quantized_np(shape, bits, out_dtype):
    min_val = 0 
    max_val = 1 << bits
    return np.random.randint(min_val, max_val, size=shape).astype(out_dtype)

with tvm.target.create('llvm'):
    #A = tvm.placeholder((batch, in_channel, in_height, in_width), dtype=input_dtype, name='A')
    #W = tvm.placeholder((num_filter, in_channel, kernel, kernel), dtype=input_dtype, name='W')
    #QW = topi.nn.bitpack(W, weight_bits, pack_axis=1, bit_axis=0, pack_type='uint8')
    
    A = tvm.placeholder((batch, in_height, in_width, in_channel), dtype=input_dtype, name='A')
    #W = tvm.placeholder((num_filter, in_channel, kernel, kernel), dtype=input_dtype, name='W')
    W = tvm.placeholder((kernel, kernel, in_channel, num_filter), dtype=input_dtype, name='W')
    
a_shape = get_const_tuple(A.shape)
w_shape = get_const_tuple(W.shape)

a_np = generate_quantized_np(a_shape, activation_bits, input_dtype)
w_np = generate_quantized_np(w_shape, weight_bits, input_dtype)

if unipolar:
    w_ = np.copy(w_np).astype(out_dtype)
    for x in np.nditer(w_, op_flags=['readwrite']):
        x[...] = 1 if x == 1 else -1
    #b_np = topi.testing.conv2d_nchw_python(a_np.astype(out_dtype), w_, stride, padding)
    b_np = topi.testing.conv2d_nhwc_python(a_np.astype(out_dtype), w_, stride, padding)
else:
    b_np = topi.testing.conv2d_nchw_python(a_np, w_np, stride, padding)
    

input_var = relay.var('input', shape=A.shape, dtype=A.dtype)
kernel_var = relay.var('kernel', shape=W.shape, dtype=W.dtype)
q_kernel = relay.nn.bitpack(kernel_var, bits=1, pack_axis=2, bit_axis=4, pack_type=input_dtype)
q_out = relay.nn.bitserial_conv2d(input_var, q_kernel, channels=32, kernel_size=(3,3), padding=(1, 1), data_layout='NHWC', pack_dtype='uint8', out_dtype='int8', kernel_layout="HWIO")

q_func = relay.Function([input_var, kernel_var], q_out)

target='aocl_sw_emu'
ctx = tvm.context(target, 0)

#target='llvm'
#ctx = tvm.cpu()
with relay.build_config(opt_level=3):
    graph, lib, params = relay.build(q_func, target=target, target_host='llvm', params={'kernel': w_np})

module = graph_runtime.create(graph, lib, ctx)
module.set_input('input', a_np)
module.set_input(**params)
module.run()

import pdb; pdb.set_trace()
output = module.get_output(0).asnumpy()
tvm.testing.assert_allclose(output, b_np, rtol=1e-5)