# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. import numpy as np import tvm from tvm import te from tvm.contrib import graph_runtime from tvm.relay.testing.config import ctx_list from tvm import relay from model_zoo import c2_squeezenet, c2_resnet50, c2_vgg19 from caffe2.python import workspace, core from caffe2.proto import caffe2_pb2 from collections import namedtuple def get_tvm_output(model, input_data, target, ctx, output_shape, output_dtype='float32'): """ Generic function to execute and get tvm output""" # supporting multiple inputs in caffe2 in a bit tricky, # because the input names can appear at the beginning or end of model.predict_net.external_input assert isinstance(input_data, np.ndarray) # here we use the first input blob to the first op to get the input name input_names = model.predict_net.op[0].input[0] shape_dict = {input_names: input_data.shape} dtype_dict = {input_names: input_data.dtype} mod, params = relay.frontend.from_caffe2( model.init_net, model.predict_net, shape_dict, dtype_dict) with relay.build_config(opt_level=3): graph, lib, params = relay.build(mod, target, params=params) m = graph_runtime.create(graph, lib, ctx) # set inputs m.set_input(input_names, tvm.nd.array(input_data.astype(input_data.dtype))) m.set_input(**params) # execute m.run() # get outputs if isinstance(output_shape, list) and isinstance(output_dtype, list): tvm_output_list = [] for i, s in enumerate(output_shape): tvm_output = m.get_output(i, tvm.nd.empty((s), output_dtype[i])) tvm_output_list.append(tvm_output.asnumpy()) return tvm_output_list else: tvm_output = m.get_output(0, tvm.nd.empty((output_shape), output_dtype)) return tvm_output.asnumpy() def get_caffe2_output(model, x, dtype='float32'): workspace.RunNetOnce(model.init_net) input_blob = model.predict_net.op[0].input[0] workspace.FeedBlob(input_blob, x.astype(dtype)) workspace.RunNetOnce(model.predict_net) output_blob = model.predict_net.external_output[0] c2_output = workspace.FetchBlob(output_blob) return c2_output def verify_caffe2_forward_impl(model, data_shape, out_shape): dtype = 'float32' data = np.random.uniform(size=data_shape).astype(dtype) c2_out = get_caffe2_output(model, data, dtype) for target, ctx in ctx_list(): tvm_out = get_tvm_output(model, data, target, ctx, out_shape, dtype) tvm.testing.assert_allclose(c2_out, tvm_out, rtol=1e-5, atol=1e-5) def test_forward_squeezenet1_1(): verify_caffe2_forward_impl(c2_squeezenet, (1, 3, 224, 224), (1, 1000, 1, 1)) def test_forward_resnet50(): verify_caffe2_forward_impl(c2_resnet50, (1, 3, 224, 224), (1, 1000)) def test_forward_vgg19(): verify_caffe2_forward_impl(c2_vgg19, (1, 3, 224, 224), (1, 1000)) Model = namedtuple('Model', ['init_net', 'predict_net']) def test_elementwise_add(): data_shape = (1, 16, 9, 9) init_net = caffe2_pb2.NetDef() init_net.name = 'test_init_net' init_net.external_output[:] = ['A', 'B'] init_net.op.extend([ core.CreateOperator( 'GivenTensorFill', [], ['A'], shape=data_shape, values=np.random.uniform(size=data_shape).flatten().tolist(), ), core.CreateOperator( 'GivenTensorFill', [], ['B'], shape=data_shape, values=np.random.uniform(size=data_shape).flatten().tolist(), ), ]) predict_net = caffe2_pb2.NetDef() predict_net.name = 'test_predict_net' predict_net.external_input[:] = ['A', 'B'] predict_net.external_output[:] = ['C'] predict_net.op.extend([ core.CreateOperator( 'Add', ['A', 'B'], ['C'], ) ]) model = Model(init_net, predict_net) verify_caffe2_forward_impl(model, data_shape, data_shape) def test_elementwise_add_with_broadcast(): data_shape = (1, 16, 9, 9) init_net = caffe2_pb2.NetDef() init_net.name = 'test_init_net' init_net.external_output[:] = ['A', 'B'] init_net.op.extend([ core.CreateOperator( 'GivenTensorFill', [], ['A'], shape=data_shape, values=np.random.uniform(size=data_shape).flatten().tolist(), ), core.CreateOperator( 'GivenTensorFill', [], ['B'], shape=(1,), values=np.random.uniform(size=1).flatten().tolist(), ), ]) predict_net = caffe2_pb2.NetDef() predict_net.name = 'test_predict_net' predict_net.external_input[:] = ['A', 'B'] predict_net.external_output[:] = ['C'] predict_net.op.extend([ core.CreateOperator( 'Add', ['A', 'B'], ['C'], broadcast=1, ) ]) model = Model(init_net, predict_net) verify_caffe2_forward_impl(model, data_shape, data_shape) def test_normalize_yuv(): data_shape = (1, 3, 96, 96) init_net = caffe2_pb2.NetDef() init_net.name = 'test_init_net' init_net.external_output[:] = ['A', 'mean', 'std'] init_net.op.extend([ core.CreateOperator( 'GivenTensorFill', [], ['A'], shape=data_shape, values=np.random.uniform(size=data_shape).flatten().tolist(), ), core.CreateOperator( 'GivenTensorFill', [], ['mean'], shape=(1, 3,), values=np.random.uniform(size=3).flatten().tolist(), ), core.CreateOperator( 'GivenTensorFill', [], ['std'], shape=(1, 3,), values=np.random.uniform(size=3).flatten().tolist(), ), ]) predict_net = caffe2_pb2.NetDef() predict_net.name = 'test_predict_net' predict_net.external_input[:] = ['A', 'mean', 'std'] predict_net.external_output[:] = ['C'] predict_net.op.extend([ core.CreateOperator( 'NormalizePlanarYUV', ['A', 'mean', 'std'], ['C'], ) ]) model = Model(init_net, predict_net) verify_caffe2_forward_impl(model, data_shape, data_shape) if __name__ == '__main__': test_forward_squeezenet1_1() test_forward_resnet50() test_forward_vgg19() test_elementwise_add() test_elementwise_add_with_broadcast() test_normalize_yuv()