| # 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 |
| import pytest |
| import tvm |
| from tvm import relay |
| from tvm.relay import transform |
| from tvm.relay.testing import run_opt_pass |
| |
| |
| def test_defuse_simple(): |
| """Simple testcase.""" |
| |
| def before(): |
| x = relay.var("x", shape=(10, 20)) |
| y = relay.add(x, relay.const(1, "float32")) |
| z = relay.exp(y) |
| w = relay.squeeze(z) |
| return relay.Function([x], w) |
| |
| x = before() |
| x = run_opt_pass(x, transform.InferType()) |
| fused = run_opt_pass(x, transform.FuseOps()) |
| defused = run_opt_pass(fused, transform.DefuseOps()) |
| |
| assert tvm.ir.structural_equal(x, defused) |
| |
| |
| def test_inception_like(): |
| def conv(data): |
| y = relay.nn.conv2d(data, relay.var("w"), kernel_size=(3, 3), padding=(1, 1), channels=16) |
| return relay.nn.relu(data=y) |
| |
| def inception_like(data): |
| c0 = conv(data) |
| c1 = conv(data) |
| return relay.concatenate((c0, c1), axis=1) |
| |
| def before(dshape): |
| x = relay.var("x", shape=dshape) |
| in1 = inception_like(x) |
| in2 = inception_like(in1) |
| return relay.Function(relay.analysis.free_vars(in2), in2) |
| |
| dshape = (1, 16, 64, 64) |
| x = before(dshape) |
| x = run_opt_pass(x, transform.InferType()) |
| fused = run_opt_pass(x, transform.FuseOps()) |
| defused = run_opt_pass(fused, transform.DefuseOps()) |
| |
| assert tvm.ir.structural_equal(x, defused) |
| |
| |
| def test_defuse_complex(): |
| """Complex defuse testcase""" |
| |
| def fused_conv2d_batch_norm(w): |
| data = relay.var("data", shape=(1, 224, 224, 3)) |
| bn_gamma0 = relay.var("bn_gamma0", relay.TensorType((64,), "float32")) |
| bn_beta0 = relay.var("bn_beta0", relay.TensorType((64,), "float32")) |
| bn_mmean0 = relay.var("bn_mean0", relay.TensorType((64,), "float32")) |
| bn_mvar0 = relay.var("bn_var0", relay.TensorType((64,), "float32")) |
| c0 = relay.nn.conv2d( |
| data, |
| w, |
| strides=(2, 2), |
| padding=(3, 3, 3, 3), |
| channels=64, |
| kernel_size=(7, 7), |
| data_layout="NHWC", |
| kernel_layout="OHWI", |
| out_layout="NHWC", |
| ) |
| c1 = relay.nn.batch_norm(c0, bn_gamma0, bn_beta0, bn_mmean0, bn_mvar0, axis=3) |
| c2 = c1[0] |
| return relay.Function(relay.analysis.free_vars(c2), c2) |
| |
| def fused_conv2d_batch_norm_relu(z): |
| data2 = relay.var("data2", shape=(1, 56, 56, 64)) |
| bn_gamma0 = relay.var("bn_gamma0", relay.TensorType((64,), "float32")) |
| bn_beta0 = relay.var("bn_beta0", relay.TensorType((64,), "float32")) |
| bn_mmean0 = relay.var("bn_mean0", relay.TensorType((64,), "float32")) |
| bn_mvar0 = relay.var("bn_var0", relay.TensorType((64,), "float32")) |
| c0 = relay.nn.conv2d( |
| data2, |
| z, |
| padding=(1, 1, 1, 1), |
| channels=64, |
| kernel_size=(3, 3), |
| data_layout="NHWC", |
| kernel_layout="OHWI", |
| out_layout="NHWC", |
| ) |
| c1 = relay.nn.batch_norm(c0, bn_gamma0, bn_beta0, bn_mmean0, bn_mvar0, axis=3) |
| c2 = c1[0] |
| c3 = relay.nn.relu(data=c2) |
| return relay.Function(relay.analysis.free_vars(c3), c3) |
| |
| def fused_max_pool2d(): |
| data1 = relay.var("data1", shape=(1, 112, 112, 64)) |
| a1 = relay.nn.max_pool2d( |
| data1, |
| pool_size=(3, 3), |
| strides=(2, 2), |
| padding=(1, 1, 1, 1), |
| layout="NHWC", |
| out_layout="NHWC", |
| ) |
| return relay.Function(relay.analysis.free_vars(a1), a1) |
| |
| def fused_add_relu(): |
| data1 = relay.var("data1", shape=(1, 56, 56, 64)) |
| data2 = relay.var("data2", shape=(1, 56, 56, 64)) |
| a0 = relay.add(data1, data2) |
| a1 = relay.nn.relu(a0) |
| return relay.Function(relay.analysis.free_vars(a1), a1) |
| |
| def before_fused(conv_layer1_weight, conv_layer2_weight): |
| data = relay.var("data", shape=(1, 3, 224, 224)) |
| data1 = relay.layout_transform(data, src_layout="NCHW", dst_layout="NHWC") |
| bn_gamma0 = relay.const(tvm.nd.array(numpy.ndarray(shape=(64,), dtype="float32"))) |
| bn_beta0 = relay.const(tvm.nd.array(numpy.ndarray(shape=(64,), dtype="float32"))) |
| bn_mmean0 = relay.const(tvm.nd.array(numpy.ndarray(shape=(64,), dtype="float32"))) |
| bn_mvar0 = relay.const(tvm.nd.array(numpy.ndarray(shape=(64,), dtype="float32"))) |
| a0 = fused_conv2d_batch_norm(conv_layer1_weight) |
| a1 = fused_max_pool2d() |
| a2 = fused_conv2d_batch_norm_relu(conv_layer2_weight) |
| a3 = fused_add_relu() |
| y0 = relay.Call(a0, [data1, bn_gamma0, bn_beta0, bn_mmean0, bn_mvar0]) |
| y1 = relay.Call(a1, [y0]) |
| y2 = relay.Call(a2, [y1, bn_gamma0, bn_beta0, bn_mmean0, bn_mvar0]) |
| y3 = relay.Call(a3, [y1, y2]) |
| return relay.Function(relay.analysis.free_vars(y3), y3) |
| |
| def golden_defused(conv_layer1_weight, conv_layer2_weight): |
| data = relay.var("data", shape=(1, 3, 224, 224)) |
| data1 = relay.layout_transform(data, src_layout="NCHW", dst_layout="NHWC") |
| bn_gamma0 = relay.const(tvm.nd.array(numpy.ndarray(shape=(64,), dtype="float32"))) |
| bn_beta0 = relay.const(tvm.nd.array(numpy.ndarray(shape=(64,), dtype="float32"))) |
| bn_mmean0 = relay.const(tvm.nd.array(numpy.ndarray(shape=(64,), dtype="float32"))) |
| bn_mvar0 = relay.const(tvm.nd.array(numpy.ndarray(shape=(64,), dtype="float32"))) |
| c0 = relay.nn.conv2d( |
| data1, |
| conv_layer1_weight, |
| strides=(2, 2), |
| padding=(3, 3, 3, 3), |
| channels=64, |
| kernel_size=(7, 7), |
| data_layout="NHWC", |
| kernel_layout="OHWI", |
| out_layout="NHWC", |
| ) |
| c1 = relay.nn.batch_norm(c0, bn_gamma0, bn_beta0, bn_mmean0, bn_mvar0, axis=3) |
| c2 = c1[0] |
| c3 = relay.nn.max_pool2d( |
| c2, |
| pool_size=(3, 3), |
| strides=(2, 2), |
| padding=(1, 1, 1, 1), |
| layout="NHWC", |
| out_layout="NHWC", |
| ) |
| c4 = relay.nn.conv2d( |
| c3, |
| conv_layer2_weight, |
| padding=(1, 1, 1, 1), |
| channels=64, |
| kernel_size=(3, 3), |
| data_layout="NHWC", |
| kernel_layout="OHWI", |
| out_layout="NHWC", |
| ) |
| c5 = relay.nn.batch_norm(c4, bn_gamma0, bn_beta0, bn_mmean0, bn_mvar0, axis=3) |
| c6 = c5[0] |
| c7 = relay.nn.relu(c6) |
| c8 = relay.add(c3, c7) |
| c9 = relay.nn.relu(c8) |
| return relay.Function(relay.analysis.free_vars(c9), c9) |
| |
| # creating weight constants for the two convolution layers |
| # in the input fused model and the golden defused model. |
| conv_layer1_weight = relay.nn.Constant( |
| tvm.nd.array(numpy.ndarray(shape=(64, 7, 7, 3), dtype="float32")) |
| ) |
| conv_layer2_weight = relay.nn.Constant( |
| tvm.nd.array(numpy.ndarray(shape=(64, 3, 3, 64), dtype="float32")) |
| ) |
| x = before_fused(conv_layer1_weight, conv_layer2_weight) |
| x = run_opt_pass(x, transform.InferType()) |
| defused = run_opt_pass(x, transform.DefuseOps()) |
| |
| golden1 = golden_defused(conv_layer1_weight, conv_layer2_weight) |
| golden1 = run_opt_pass(golden1, transform.InferType()) |
| |
| assert tvm.ir.structural_equal(defused, golden1), ( |
| "Actual = \n" + str(defused) + "\nGolden = \n" + str(golden1) |
| ) |
| |
| |
| if __name__ == "__main__": |
| tvm.testing.main() |
