(device, batch_size, channels, use_trace, symbolic, gopt_level, dtype)
| 63 | @pytest.mark.parametrize("gopt_level", [None, 1, 2]) |
| 64 | @pytest.mark.parametrize("dtype", ["float32"]) |
| 65 | def test_subgraph(device, batch_size, channels, use_trace, symbolic, gopt_level, dtype): |
| 66 | device = CompNode(device) |
| 67 | |
| 68 | def subgraph_batch_norm(inp, weight, bias, eps, diff): |
| 69 | inp = inp.detach() |
| 70 | with GradManager().attach(inp) as gm: |
| 71 | batch_norm_fn = _get_batch_norm_fn( |
| 72 | dtype, device, channels, ndim, interpret=False, gopt_level=gopt_level |
| 73 | ) |
| 74 | out, *_ = batch_norm_fn(inp, eps, weight, bias) |
| 75 | gm.backward(out * 1e3 + 1e3, diff) |
| 76 | return out, inp.grad |
| 77 | |
| 78 | def primitive_batch_norm(inp, weight, bias, eps, diff): |
| 79 | inp = inp.detach() |
| 80 | with GradManager().attach(inp) as gm: |
| 81 | batch_norm_fn = _get_batch_norm_fn( |
| 82 | dtype, device, channels, ndim, interpret=True, gopt_level=gopt_level |
| 83 | ) |
| 84 | (out,) = batch_norm_fn(inp, eps, weight, bias) |
| 85 | gm.backward(out * 1e3 + 1e3, diff) |
| 86 | return out, inp.grad |
| 87 | |
| 88 | if use_trace: |
| 89 | subgraph_batch_norm = trace(symbolic=symbolic)(subgraph_batch_norm) |
| 90 | primitive_batch_norm = trace(symbolic=symbolic)(primitive_batch_norm) |
| 91 | |
| 92 | def rand_tensor(shape, dtype=dtype, device=device): |
| 93 | return megengine.tensor(np.random.random(shape), dtype=dtype, device=device) |
| 94 | |
| 95 | # skip this test because could not do several reduce sequentially with opr cache |
| 96 | return |
| 97 | |
| 98 | # test shape change |
| 99 | for image_shape in [(223, 223), (10, 20)]: |
| 100 | ndim = len(image_shape) + 2 |
| 101 | input_shape = (batch_size, channels) + image_shape |
| 102 | param_shape = (1, channels) + (1,) * len(image_shape) |
| 103 | |
| 104 | inp = rand_tensor(input_shape) * 1e3 + 1e3 |
| 105 | weight = rand_tensor(param_shape) |
| 106 | bias = rand_tensor(param_shape) |
| 107 | eps = megengine.tensor(1e-5, dtype=dtype, device=device) |
| 108 | |
| 109 | diff = rand_tensor(input_shape) |
| 110 | |
| 111 | out1, grad1 = subgraph_batch_norm(inp, weight, bias, eps, diff) |
| 112 | out2, grad2 = primitive_batch_norm(inp, weight, bias, eps, diff) |
| 113 | |
| 114 | _assert_allclose(out1.numpy(), out2.numpy()) |
| 115 | _assert_allclose(grad1.numpy(), grad2.numpy()) |
| 116 | |
| 117 | |
| 118 | @functools.lru_cache(maxsize=None) |
nothing calls this directly
no test coverage detected