test_autodiff.py 8.2 KB
Newer Older
1 2 3 4 5 6 7 8
# -*- coding: utf-8 -*-
# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
#
# Copyright (c) 2014-2020 Megvii Inc. All rights reserved.
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
9
import gc
10
import platform
11 12 13 14 15 16 17
import weakref

import numpy as np
import pytest

import megengine as mge
import megengine.distributed as dist
18
import megengine.functional as F
19 20 21
from megengine.core._imperative_rt import TensorAttr, core2, imperative
from megengine.core._imperative_rt.core2 import TensorWeakRef, apply
from megengine.core._imperative_rt.imperative import sync
22 23
from megengine.core.autodiff.grad import Grad
from megengine.core.ops.builtin import Elemwise
24
from megengine.distributed.helper import get_device_count_by_fork
25 26 27 28
from megengine.functional.distributed import remote_recv, remote_send


def _elwise(mode):
29
    op = Elemwise(mode)
30 31 32 33 34 35 36 37

    def f(*args):
        (result,) = apply(op, *args)
        return result

    return f


38 39 40 41
add = _elwise(Elemwise.Mode.ADD)
mul = _elwise(Elemwise.Mode.MUL)
cos = _elwise(Elemwise.Mode.COS)
relu = _elwise(Elemwise.Mode.RELU)
42 43 44


def as_tensor(x):
45
    return mge.Tensor(x)
46 47 48


def save_to(self, name="grad"):
49
    def callback(grad):
50 51 52 53 54
        setattr(self, name, grad)

    return callback


55 56 57
@pytest.mark.skipif(
    platform.system() == "Darwin", reason="do not imp GPU mode at macos now"
)
58 59 60
@pytest.mark.skipif(
    platform.system() == "Windows", reason="windows disable MGB_ENABLE_OPR_MM"
)
61 62
@pytest.mark.skipif(get_device_count_by_fork("gpu") < 2, reason="need more gpu device")
@pytest.mark.isolated_distributed
63 64 65
def test_dist_grad():
    world_size = 2
    x_np = np.random.rand(10).astype("float32")
66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90

    @dist.launcher
    def worker():
        rank = dist.get_rank()
        if rank == 0:
            grad = Grad()

            x = as_tensor(x_np)
            grad.wrt(x, callback=save_to(x))
            # need a placeholder to trace operator
            send_x = remote_send(x, 1)
            recv_x = remote_recv(1, x_np.shape, x_np.dtype)
            y = recv_x * recv_x

            grad([y], [as_tensor(np.ones_like(x_np))])
            np.testing.assert_almost_equal(x.grad.numpy(), x.numpy() * 2)
        elif rank == 1:
            grad = Grad()

            recv_x = remote_recv(0, x_np.shape, x_np.dtype)
            send_x = remote_send(recv_x, 0)

            grad([], [])

    worker()
91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137


def test_grad():
    x_np = np.random.rand(10).astype("float32")
    x = as_tensor(x_np)

    grad = Grad().wrt(x, callback=save_to(x))

    y = cos(x)

    grad(y, as_tensor(np.ones_like(x_np)))
    np.testing.assert_almost_equal(x.grad.numpy(), -np.sin(x_np))


def test_grad_2():
    x_np = np.random.rand(10).astype("float32")
    x = as_tensor(x_np)

    grad = Grad().wrt(x, callback=save_to(x))

    y = mul(x, x)
    y = mul(y, y)

    grad(y, as_tensor(np.ones_like(x_np)))
    np.testing.assert_almost_equal(x.grad.numpy(), 4 * x_np ** 3, decimal=6)


@pytest.mark.skip(reason="high order gradient was not implemented yet")
def test_2nd_grad():
    x_np = np.random.rand(10).astype("float32")
    x = as_tensor(x_np)
    ones = as_tensor(np.ones_like(x_np))

    grad = Grad().wrt(x, callback=save_to(x))
    grad2 = Grad().wrt(x, callback=save_to(x))

    y = cos(x)

    grad(y, ones)
    np.testing.assert_almost_equal(x.grad.numpy(), -np.sin(x_np), decimal=5)

    grad2(x.grad, ones)
    np.testing.assert_almost_equal(x.grad.numpy(), -np.cos(x_np))


def test_grad_with_tensor_wrapper():
    x_np = np.random.rand(10).astype("float32")
138
    x = mge.Tensor(x_np)
139 140 141 142 143 144

    grad = Grad().wrt(x, callback=save_to(x))

    y = mul(x, x)
    y = mul(y, y)

145
    grad(y, mge.Tensor(np.ones_like(x_np)))
146 147 148
    np.testing.assert_almost_equal(x.grad.numpy(), 4 * x_np ** 3, decimal=6)


149 150 151 152
def test_release():
    def check(f):
        n = 0
        d = None
153 154 155 156 157 158 159 160 161 162
        gc.disable()
        try:
            for i in range(3):
                f()
                m = len(gc.get_objects())
                d = m - n
                n = m
            assert d == 0
        finally:
            gc.enable()
163

164 165
    x = mge.Tensor([0.0])
    dy = mge.Tensor(np.ones_like(x.numpy()))
166 167 168 169 170 171 172 173 174

    @check
    def _():
        g = Grad().wrt(x)
        y = x * x
        g(y, dy)

    @check
    def _():
175
        with Grad().wrt(x):
176 177 178 179
            pass

    @check
    def _():
180
        with Grad().wrt(x):
181 182 183
            y = x * x


184 185
def test_grad_inplace():
    x_np = np.random.rand(10).astype("float32")
186
    x = mge.Tensor(x_np)
187 188 189 190 191 192

    grad = Grad().wrt(x, callback=save_to(x))

    y = mul(x, x)
    y *= y

193
    grad(y, mge.Tensor(np.ones_like(x_np)))
194 195 196 197 198 199 200
    np.testing.assert_almost_equal(x.grad.numpy(), 4 * x_np ** 3, decimal=6)


def test_elemwise_add():
    x_np = np.random.rand(10).astype("float32")
    y_np = np.random.rand(10, 10).astype("float32")
    dz_np = np.random.rand(10, 10).astype("float32")
201 202 203
    x = mge.Tensor(x_np)
    y = mge.Tensor(y_np)
    dz = mge.Tensor(dz_np)
204 205 206 207 208

    refs = {}

    def f(x, y):
        x = x * 2
209 210
        refs["x"] = TensorWeakRef(x)
        refs["y"] = TensorWeakRef(y)
211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227
        return x + y

    grad = Grad().wrt(x, callback=save_to(x))

    z = f(x, y)
    del y

    for k, r in refs.items():
        assert r() is None

    grad(z, dz)
    np.testing.assert_almost_equal(x.grad.numpy(), dz_np.sum(0) * 2, decimal=5)


def test_elemwise_relu():
    x_np = [1.0, -1.0]
    dz_np = [1.0]
228 229
    x = mge.Tensor(x_np)
    dz = mge.Tensor(dz_np)
230 231 232 233 234

    refs = {}

    def f(x):
        x = x * 2
235
        refs["x"] = TensorWeakRef(x)
236 237 238 239 240 241 242 243 244 245 246 247 248
        return relu(x)

    grad = Grad().wrt(x, callback=save_to(x))

    z = f(x)

    assert refs["x"]() is None

    grad(z, dz)
    np.testing.assert_almost_equal(x.grad.numpy(), [2.0, 0])


def test_elemwise_relu_backward_fn():
249
    op = Elemwise(Elemwise.Mode.RELU)
250 251 252 253 254 255
    attr = TensorAttr()
    attr.dtype = "float32"
    attr.comp_node = "xpux"
    result = imperative.make_backward_graph(op, [attr], [True], [True])
    backward_graph, save_for_backward_mask, input_has_grad = result
    assert save_for_backward_mask == [False, True, True], save_for_backward_mask
256 257 258 259


def test_reshape():
    x_np = np.random.rand(2, 5).astype("float32")
260
    x = mge.Tensor(x_np)
261 262 263 264 265 266 267 268 269 270

    grad = Grad().wrt(x, callback=save_to(x))
    y = x.reshape(5, 2)

    grad(y, F.ones_like(y))
    np.testing.assert_equal(np.ones((2, 5), dtype=np.float32), x.grad.numpy())


def test_subtensor():
    x_np = np.random.rand(3, 3).astype("float32")
271
    x = mge.Tensor(x_np)
272 273 274 275 276 277 278 279 280 281 282 283

    grad = Grad().wrt(x, callback=save_to(x))
    y = x[1:-1, :2]

    grad(y, F.ones_like(y))
    np.testing.assert_equal(
        np.array([[0, 0, 0], [1, 1, 0], [0, 0, 0]], dtype=np.float32), x.grad.numpy()
    )


def test_IndexingMultiAxisVec():
    x_np = np.random.rand(3, 3).astype("float32")
284
    x = mge.Tensor(x_np)
285 286 287 288 289 290 291 292 293 294 295 296

    grad = Grad().wrt(x, callback=save_to(x))
    y = x[[0, 2], [0, 2]]

    grad(y, F.ones_like(y))
    np.testing.assert_equal(
        np.array([[1, 0, 0], [0, 0, 0], [0, 0, 1]], dtype=np.float32), x.grad.numpy()
    )


def test_AxisAddRemove():
    x_np = np.random.rand(1, 5).astype("float32")
297
    x = mge.Tensor(x_np)
298 299

    grad = Grad().wrt(x, callback=save_to(x))
300
    y = F.squeeze(F.expand_dims(x, 2), 0)
301 302 303 304 305 306 307 308 309

    grad(y, F.ones_like(y))
    np.testing.assert_equal(
        np.array([[1, 1, 1, 1, 1]], dtype=np.float32), x.grad.numpy()
    )


def test_Broadcast():
    x_np = np.random.rand(3, 3, 1).astype("float32")
310
    x = mge.Tensor(x_np)
311 312

    grad = Grad().wrt(x, callback=save_to(x))
313
    y = F.broadcast_to(x, (3, 3, 10))
314 315 316 317 318 319 320

    grad(y, F.ones_like(y))
    np.testing.assert_equal(np.ones((3, 3, 1), dtype=np.float32) * 10, x.grad.numpy())


def test_Reduce_sum():
    x_np = np.random.rand(3, 3).astype("float32")
321
    x = mge.Tensor(x_np)
322 323 324 325 326 327 328 329 330 331

    grad = Grad().wrt(x, callback=save_to(x))
    y = x.sum(axis=0)

    grad(y, F.ones_like(y))
    np.testing.assert_equal(np.ones((3, 3), dtype=np.float32), x.grad.numpy())


def test_Reduce_mean():
    x_np = np.random.rand(3, 3).astype("float32")
332
    x = mge.Tensor(x_np)
333 334 335 336 337 338

    grad = Grad().wrt(x, callback=save_to(x))
    y = x.mean(axis=0)

    grad(y, F.ones_like(y))
    np.testing.assert_equal(np.ones((3, 3), dtype=np.float32) / 3, x.grad.numpy())