test_op_grad_level1.py 4.87 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
# 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.
17
import numpy as np
18
import pytest
SWu committed
19

Zhi committed
20
import tvm
21
from tvm import te
22
from tvm import relay
SWu committed
23
from tvm.relay.testing import check_grad, ctx_list, run_infer_type
Zhi committed
24
from tvm.relay.transform import gradient
SWu committed
25

Zhi committed
26

27 28 29 30
def sigmoid(x):
    one = np.ones_like(x)
    return one / (one + np.exp(-x))

Zhi committed
31

32 33 34 35 36
def relu(x):
    x_copy = np.copy(x)
    np.maximum(x_copy, 0, x_copy)
    return x_copy

SWu committed
37

38 39 40 41 42 43 44 45 46 47 48 49
def test_unary_op():
    def check_single_op(opfunc, ref):
        shape = (10, 4)
        dtype = 'float32'
        tp = relay.TensorType(shape, dtype)
        x = relay.var("x", tp)
        y = opfunc(x)

        if ref is not None:
            data = np.random.rand(*shape).astype(dtype)
            ref_grad = ref(data)
            fwd_func = relay.Function([x], y)
50
            fwd_func = run_infer_type(fwd_func)
Zhi committed
51
            bwd_func = run_infer_type(gradient(fwd_func))
52 53 54 55 56 57 58 59 60 61 62

            for target, ctx in ctx_list():
                intrp = relay.create_executor(ctx=ctx, target=target)
                op_res, (op_grad, ) = intrp.evaluate(bwd_func)(data)
                np.testing.assert_allclose(op_grad.asnumpy(), ref_grad, rtol=0.01)

    for opfunc, ref in [(tvm.relay.log, lambda x: 1 / x),
                        (tvm.relay.exp, np.exp),
                        (tvm.relay.sigmoid, lambda x: sigmoid(x) * (1 - sigmoid(x))),
                        (tvm.relay.tanh, lambda x: 1 - np.tanh(x) * np.tanh(x)),
                        (tvm.relay.sqrt, lambda x: 0.5 * np.power(x, -0.5)),
63
                        (tvm.relay.abs, lambda x: np.where(x < 0, -np.ones_like(x), np.ones_like(x))),
64 65
                        (relay.nn.relu, lambda x: np.where(x < 0, np.zeros_like(x), np.ones_like(x))),
                        (tvm.relay.cos, lambda x: -1.0 * np.sin(x)),
66
                        (tvm.relay.sin, lambda x: np.cos(x)),
67
                        (tvm.relay.tan, lambda x: 1.0 / (np.cos(x) ** 2)),
68
                        (tvm.relay.atan, lambda x: 1 / (1 + np.power(x, 2.0)))]:
69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86
        check_single_op(opfunc, ref)


def test_binary_op():
    def inst(vars, sh):
        return [vars.get(s, s) for s in sh]

    def check_binary_op(opfunc, ref):
        s = (5, 10, 5)
        t = relay.TensorType((5, 10, 5))
        x = relay.var("x", t)
        y = relay.var("y", t)
        z = opfunc(x, y)

        x_data = np.random.rand(*s).astype(t.dtype)
        y_data = np.random.rand(*s).astype(t.dtype)
        ref_grad0, ref_grad1 = ref(x_data, y_data)
        fwd_func = relay.Function([x, y], z)
87
        fwd_func = run_infer_type(fwd_func)
Zhi committed
88
        bwd_func = run_infer_type(gradient(fwd_func))
89 90 91 92 93 94 95 96 97 98 99 100 101 102

        for target, ctx in ctx_list():
            intrp = relay.create_executor(ctx=ctx, target=target)
            op_res, (op_grad0, op_grad1) = intrp.evaluate(bwd_func)(x_data, y_data)
            np.testing.assert_allclose(op_grad0.asnumpy(), ref_grad0, rtol=0.01)
            np.testing.assert_allclose(op_grad1.asnumpy(), ref_grad1, rtol=0.01)

    for opfunc, ref in [(relay.add, lambda x, y: [np.ones_like(x), np.ones_like(y)]),
                        (relay.subtract, lambda x, y: [np.ones_like(x), -np.ones_like(y)]),
                        (relay.multiply, lambda x, y: [y, x]),
                        (relay.divide, lambda x, y: [1 / y, - x / (y**2)])]:
        check_binary_op(opfunc, ref)


SWu committed
103 104 105
def test_softmax_grad():
    data = relay.var("data", relay.TensorType((1, 16), "float64"))
    fwd_func = relay.Function([data], relay.nn.softmax(data))
106 107 108 109 110 111 112
    check_grad(fwd_func, scale=1)


def test_log_softmax_grad():
    data = relay.var("data", relay.TensorType((2, 16), "float64"))
    fwd_func = relay.Function([data], relay.nn.log_softmax(data))
    check_grad(fwd_func, scale=1)
SWu committed
113 114


SWu committed
115 116 117 118
def verify_bias_add(d_shape, b_shape, axis=1):
    data = relay.var("data", relay.TensorType(d_shape, "float32"))
    bias = relay.var("bias", relay.TensorType(b_shape, "float32"))
    fwd_func = relay.Function([data, bias], relay.nn.bias_add(data, bias, axis=axis))
SWu committed
119 120 121
    check_grad(fwd_func)


SWu committed
122 123 124 125 126 127 128
def test_bias_add_grad():
    verify_bias_add((1, 16), (16,))
    verify_bias_add((1, 8, 2, 2), (8,))
    verify_bias_add((1, 2, 2, 8), (8,), 3)
    verify_bias_add((4, 8), (8,))


129
if __name__ == "__main__":
130
    pytest.main([__file__])