| 20 | @pytest.mark.isolated_distributed |
| 21 | @pytest.mark.parametrize("enable_amp", [False, True]) |
| 22 | def test_syncbn(enable_amp): |
| 23 | nr_chan = 8 |
| 24 | data_shape = (3, nr_chan, 4, 16) |
| 25 | momentum = 0.9 |
| 26 | eps = 1e-5 |
| 27 | running_mean = np.zeros((1, nr_chan, 1, 1), dtype=np.float32) |
| 28 | running_var = np.ones((1, nr_chan, 1, 1), dtype=np.float32) |
| 29 | steps = 4 |
| 30 | nr_ranks = 2 |
| 31 | server = dist.Server() |
| 32 | port = server.py_server_port |
| 33 | |
| 34 | @dist.launcher(n_gpus=2) |
| 35 | def worker(data, yv_expect, running_mean, running_var): |
| 36 | with amp.autocast(enabled=enable_amp): |
| 37 | rank = dist.get_rank() |
| 38 | bn = SyncBatchNorm(nr_chan, momentum=momentum, eps=eps) |
| 39 | for i in range(steps): |
| 40 | yv = bn(Tensor(data[rank][i])) |
| 41 | if enable_amp: |
| 42 | np.testing.assert_allclose( |
| 43 | yv.numpy(), yv_expect[rank], atol=5e-4, rtol=5e-4 |
| 44 | ) |
| 45 | else: |
| 46 | _assert_allclose(yv.numpy(), yv_expect[rank]) |
| 47 | _assert_allclose(bn.running_mean.numpy(), running_mean) |
| 48 | _assert_allclose(bn.running_var.numpy(), running_var) |
| 49 | |
| 50 | xv = [] |
| 51 | for i in range(steps): |
| 52 | xv.append(np.random.normal(loc=2.3, size=data_shape).astype(np.float32)) |
| 53 | xv_transposed = np.transpose(xv[i], [0, 2, 3, 1]).reshape( |
| 54 | (data_shape[0] * data_shape[2] * data_shape[3], nr_chan) |
| 55 | ) |
| 56 | |
| 57 | mean = np.mean(xv_transposed, axis=0).reshape(1, nr_chan, 1, 1) |
| 58 | |
| 59 | var_biased = np.var(xv_transposed, axis=0).reshape((1, nr_chan, 1, 1)) |
| 60 | sd = np.sqrt(var_biased + eps) |
| 61 | |
| 62 | var_unbiased = np.var(xv_transposed, axis=0, ddof=1).reshape((1, nr_chan, 1, 1)) |
| 63 | running_mean = running_mean * momentum + mean * (1 - momentum) |
| 64 | running_var = running_var * momentum + var_unbiased * (1 - momentum) |
| 65 | |
| 66 | yv_expect = (xv[i] - mean) / sd |
| 67 | |
| 68 | data = [] |
| 69 | for i in range(nr_ranks): |
| 70 | data.append([]) |
| 71 | for j in range(steps): |
| 72 | data[i].append(xv[j][:, :, :, i * 8 : i * 8 + 8]) |
| 73 | |
| 74 | yv_expect = [yv_expect[:, :, :, i * 8 : i * 8 + 8] for i in range(nr_ranks)] |
| 75 | |
| 76 | worker(data, yv_expect, running_mean, running_var) |
| 77 | |
| 78 | |
| 79 | def test_batchnorm(): |