test_fold_axis.py

"""Unittest cases for fold_axis"""
import tvm
import nnvm
import nnvm.testing.resnet
import numpy as np
from nnvm import symbol as sym
from nnvm.compiler import graph_util, graph_attr

def test_fold_axis_conv():
    # Before simplify
    def before(x, conv_weight, conv_bias, in_scale, out_scale, channels):
        x = x * sym.expand_dims(in_scale, axis=1, num_newaxis=2)
        y = sym.conv2d(x, conv_weight, conv_bias,
                       channels=channels,
                       kernel_size=(3, 3),
                       padding=(1, 1),
                       name="conv")
        y = sym.relu(y)
        y = y * sym.expand_dims(out_scale, axis=1, num_newaxis=2)
        return y

    def expected(x, conv_weight, conv_bias, in_scale, out_scale, channels):
        conv_weight = conv_weight * sym.expand_dims(out_scale, axis=1, num_newaxis=3)
        conv_weight = conv_weight * sym.expand_dims(in_scale, axis=1, num_newaxis=2)
        conv_bias = conv_bias * out_scale
        y = sym.conv2d(x,
                       conv_weight,
                       conv_bias,
                       channels=channels,
                       kernel_size=(3, 3),
                       padding=(1, 1),
                       name="conv")
        y = sym.relu(y)
        return y

    def check(shape, channels):
        x = sym.Variable("x") + 1
        weight = sym.Variable("weight")
        bias = sym.Variable("bias")
        in_scale = sym.Variable("in_scale")
        out_scale = sym.Variable("out_scale")
        y1 = before(x, weight, bias, in_scale, out_scale, channels)
        y2 = expected(x, weight, bias, in_scale, out_scale, channels)
        ishape = {"x": shape, "out_scale": (channels,), "in_scale": (shape[1],)}
        g1 = nnvm.graph.create(y1)
        g2 = nnvm.graph.create(y2)
        graph_attr.set_shape_inputs(g1, ishape)
        g1 = g1.apply("InferShape").apply("FoldScaleAxis")
        # assert graph equals as expected
        graph_util.check_graph_equal(g1, g2)

    check((2, 4, 10, 10), 2)

def test_fold_axis_depthwise_conv():
    # Before simplify
    def before(x, conv_weight, conv_bias, in_scale, out_scale, channels):
        x = x * sym.expand_dims(in_scale, axis=1, num_newaxis=2)
        y = sym.conv2d(x, conv_weight, conv_bias,
                       channels=channels,
                       kernel_size=(3, 3),
                       padding=(1, 1),
                       groups=54,
                       name="depthiwise_conv")
        y = sym.relu(y)
        y = y * sym.expand_dims(out_scale, axis=1, num_newaxis=2)
        return y

    def expected(x, conv_weight, conv_bias, in_scale, out_scale, channels):
        conv_weight = conv_weight * sym.expand_dims(out_scale, axis=1, num_newaxis=3)
        conv_weight = conv_weight * sym.expand_dims(in_scale, axis=1, num_newaxis=3)
        conv_bias = conv_bias * out_scale
        y = sym.conv2d(x,
                       conv_weight,
                       conv_bias,
                       channels=channels,
                       kernel_size=(3, 3),
                       padding=(1, 1),
                       groups=54,
                       name="depthiwise_conv")
        y = sym.relu(y)
        return y

    def check(shape, channels):
        x = sym.Variable("x") + 1
        weight = sym.Variable("weight")
        bias = sym.Variable("bias")
        in_scale = sym.Variable("in_scale")
        out_scale = sym.Variable("out_scale")
        y1 = before(x, weight, bias, in_scale, out_scale, channels)
        y2 = expected(x, weight, bias, in_scale, out_scale, channels)
        ishape = {"x": shape, "out_scale": (channels,), "in_scale": (shape[1],)}
        g1 = nnvm.graph.create(y1)
        g2 = nnvm.graph.create(y2)
        graph_attr.set_shape_inputs(g1, ishape)
        g1 = g1.apply("InferShape").apply("FoldScaleAxis")
        # assert graph equals as expected
        graph_util.check_graph_equal(g1, g2)

    check((1, 54, 63, 127), 54)

def test_fold_fail():
    # Before simplify
    def before(x, scale, channels):
        y = sym.conv2d(x,
                       channels=channels,
                       kernel_size=(3, 3),
                       padding=(1, 1),
                       name="conv")
        y = y * sym.expand_dims(scale, axis=1, num_newaxis=1)
        return y

    def check(shape, channels):
        x = sym.Variable("x")
        bias = sym.Variable("bias")
        scale = sym.Variable("scale")
        y1 = before(x, scale, channels)
        ishape = {"x": shape, "scale": (channels,), "bias": (channels,)}
        g1 = nnvm.graph.create(y1)
        graph_attr.set_shape_inputs(g1, ishape)
        g2 = g1.apply("InferShape").apply("FoldScaleAxis")
        # assert graph equals as expected
        graph_util.check_graph_equal(g1, g2)

    check((2, 10, 10, 10), 10)


def test_fold_resnet():
    batch_size = 1
    num_classes = 1000
    image_shape = (3, 224, 224)
    data_shape = (batch_size,) +image_shape
    net, params = nnvm.testing.resnet.get_workload(
        batch_size=1, image_shape=image_shape)
    ishape = {"data" : data_shape}
    graph = nnvm.graph.create(net)
    data = np.random.uniform(size=data_shape).astype("float32")
    # Initial pass do shape type inference
    shape, _ = graph_util.infer_shape(graph, **ishape)
    ishape.update(zip(graph.index.input_names, shape))

    def run_prune(graph, params, opt_level):
        # Apply optimization
        with nnvm.compiler.build_config(opt_level=0):
            graph = nnvm.compiler.optimize(graph, ishape)
        graph, params = nnvm.compiler.build_module.precompute_prune(graph, params)
        params["data"] = data
        return nnvm.compiler.build_module._run_graph(graph, params)

    x = run_prune(graph, params, 0)
    y = run_prune(graph, params, 3)
    tvm.testing.assert_allclose(y[0].asnumpy(), x[0].asnumpy())


if __name__ == "__main__":
    test_fold_resnet()
    test_fold_axis_conv()
    test_fold_fail()
    test_fold_axis_depthwise_conv()