def test_adt_list(): mod = relay.Module() p = Prelude(mod) l1 = p.cons(relay.const(1), p.nil()) l21 = p.cons(relay.const(2), l1) l321 = p.cons(relay.const(3), l21) f = relay.Function([], l321) mod["main"] = f vm = create_vm(mod) ser = serializer.Serializer(vm) code, lib = ser.serialize() deser = deserializer.Deserializer(code, lib) des_vm = deser.deserialize() result = veval(des_vm) assert len(result) == 2 assert len(result[1]) == 2 assert len(result[1][1]) == 2 res = [] res.append(result[0].asnumpy().tolist()) res.append(result[1][0].asnumpy().tolist()) res.append(result[1][1][0].asnumpy().tolist()) tvm.testing.assert_allclose(res, np.array([3, 2, 1]))
def test_loop(): mod = relay.module.Module({}) sum_up = relay.GlobalVar('sum_up') i = relay.var('i', shape=[], dtype='int32') accum = relay.var('accum', shape=[], dtype='int32') sb = ScopeBuilder() with sb.if_scope(relay.equal(i, relay.const(0, 'int32'))): sb.ret(accum) with sb.else_scope(): one_less = relay.subtract(i, relay.const(1, 'int32')) new_accum = relay.add(accum, i) sb.ret(relay.Call(sum_up, [one_less, new_accum])) func = relay.Function([i, accum], sb.get()) mod[sum_up] = func loop_bound = 0 i_data = np.array(loop_bound, dtype='int32') accum_data = np.array(0, dtype='int32') iarg = relay.var('i', shape=[], dtype='int32') aarg = relay.var('accum', shape=[], dtype='int32') mod["main"] = relay.Function([iarg, aarg], sum_up(iarg, aarg)) vm = create_vm(mod) ser = serializer.Serializer(vm) code, lib = ser.serialize() deser = deserializer.Deserializer(code, lib) des_vm = deser.deserialize() result = veval(des_vm, i_data, accum_data) tvm.testing.assert_allclose(result.asnumpy(), sum(range(1, loop_bound + 1)))
def test_save_load(): x = relay.var('x', shape=(10, 10)) f = relay.Function([x], x + x) x_data = np.random.rand(10, 10).astype('float32') # serialize. vm = create_vm(f) ser = serializer.Serializer(vm) code, lib = ser.serialize() assert isinstance(code, bytearray) # save and load the code and lib file. tmp = util.tempdir() path_lib = tmp.relpath("lib.so") lib.export_library(path_lib) with open(tmp.relpath("code.bc"), "wb") as fo: fo.write(code) loaded_lib = tvm.module.load(path_lib) loaded_code = bytearray(open(tmp.relpath("code.bc"), "rb").read()) # deserialize. deser = deserializer.Deserializer(loaded_code, loaded_lib) des_vm = deser.deserialize() res = veval(des_vm, x_data) tvm.testing.assert_allclose(res.asnumpy(), x_data + x_data)
def get_serialized_output(mod, data, params, target, ctx, dtype='float32'): vm = create_vm(mod, ctx, target) ser = serializer.Serializer(vm) code, lib = ser.serialize() deser = deserializer.Deserializer(code, lib) des_vm = deser.deserialize() des_vm.init(ctx) des_vm.load_params(params) result = des_vm.run(data) return result.asnumpy().astype(dtype)
def test_const(): c = relay.const(1.0, "float32") x = relay.var('x', shape=(10, 10), dtype='float32') f = relay.Function([x], x + c) vm = create_vm(f) ser = serializer.Serializer(vm) code, lib = ser.serialize() assert isinstance(code, bytearray) deser = deserializer.Deserializer(code, lib) des_vm = deser.deserialize() x_data = np.random.rand(10, 10).astype('float32') res = veval(des_vm, x_data) tvm.testing.assert_allclose(res.asnumpy(), x_data + 1)
def test_tuple(): ttype = relay.TupleType([relay.TensorType((1,)), relay.TensorType((10,))]) tup = relay.var('tup', type_annotation=ttype) f = relay.Function([tup], relay.TupleGetItem(tup, 1)) i_data = np.random.rand(41).astype('float32') j_data = np.random.rand(10).astype('float32') vm = create_vm(f) ser = serializer.Serializer(vm) code, lib = ser.serialize() deser = deserializer.Deserializer(code, lib) des_vm = deser.deserialize() result = veval(des_vm, (i_data, j_data)) tvm.testing.assert_allclose(result.asnumpy(), j_data)
def test_closure(): x = relay.var('x', shape=()) y = relay.var('y', shape=()) f = relay.Function([x], x + y) ff = relay.Function([y], f) clo = ff(relay.const(1.0)) main = clo(relay.const(2.0)) vm = create_vm(main) ser = serializer.Serializer(vm) code, lib = ser.serialize() deser = deserializer.Deserializer(code, lib) des_vm = deser.deserialize() res = veval(des_vm) tvm.testing.assert_allclose(res.asnumpy(), 3.0)
def test_adt_compose(): mod = relay.Module() p = Prelude(mod) compose = p.compose # add_one = fun x -> x + 1 sb = relay.ScopeBuilder() x = relay.var('x', 'float32') x1 = sb.let('x1', x) xplusone = x1 + relay.const(1.0, 'float32') sb.ret(xplusone) body = sb.get() add_one = relay.GlobalVar("add_one") add_one_func = relay.Function([x], body) # add_two = compose(add_one, add_one) sb = relay.ScopeBuilder() y = relay.var('y', 'float32') add_two_func = sb.let('add_two', compose(add_one_func, add_one_func)) add_two_res = add_two_func(y) sb.ret(add_two_res) add_two_body = sb.get() mod[add_one] = add_one_func f = relay.Function([y], add_two_body) mod["main"] = f vm = create_vm(mod) ser = serializer.Serializer(vm) code, lib = ser.serialize() deser = deserializer.Deserializer(code, lib) des_vm = deser.deserialize() x_data = np.array(np.random.rand()).astype('float32') result = veval(des_vm, x_data) tvm.testing.assert_allclose(result.asnumpy(), x_data + 2.0)
def test_if(): x = relay.var('x', shape=(10, 10)) y = relay.var('y', shape=(10, 10)) equal = relay.op.equal(x, y) equal = relay.op.nn.batch_flatten(equal) f = relay.Function([x, y], relay.If(relay.op.min(equal, axis=[0, 1]), x, y)) x_data = np.random.rand(10, 10).astype('float32') y_data = np.random.rand(10, 10).astype('float32') vm = create_vm(f) ser = serializer.Serializer(vm) code, lib = ser.serialize() deser = deserializer.Deserializer(code, lib) des_vm = deser.deserialize() # same res = veval(des_vm, x_data, x_data) tvm.testing.assert_allclose(res.asnumpy(), x_data) # diff res = veval(des_vm, x_data, y_data) tvm.testing.assert_allclose(res.asnumpy(), y_data)