def __init__(self):
super().__init__()
self.bn = BatchNorm2d(4)
self.seq = Sequential(
BatchNorm2d(4),
self.InnerModule(),
)
def __init__(self):
super().__init__()
self.bn = BatchNorm2d(4)
self.a = [
BatchNorm2d(4),
{"x": BatchNorm2d(4), "y": [BatchNorm2d(4), self.InnerModule()], "z": 0},
(self.InnerModule(),),
]
def __init__(self, has_bn=False):
super().__init__()
self.conv0 = Conv2d(1, 20, kernel_size=5, bias=True)
self.pool0 = AvgPool2d(2)
self.conv1 = Conv2d(20, 20, kernel_size=5, bias=True)
self.pool1 = AvgPool2d(2)
self.fc0 = Linear(20 * 4 * 4, 500, bias=True)
self.fc1 = Linear(500, 10, bias=True)
self.bn0 = None
self.bn1 = None
if has_bn:
self.bn0 = BatchNorm2d(20)
self.bn1 = BatchNorm2d(20)
def __init__(self):
super().__init__()
self.conv1 = Conv2d(3, 128, 3, padding=1, bias=False)
self.conv2 = Conv2d(3, 128, 3, dilation=2, bias=False)
self.bn1 = BatchNorm1d(128)
self.bn2 = BatchNorm2d(128)
self.pooling = MaxPool2d(kernel_size=2, padding=0)
modules = OrderedDict()
modules["depthwise"] = Conv2d(
256,
256,
3,
1,
0,
groups=256,
bias=False,
)
modules["pointwise"] = Conv2d(
256,
256,
kernel_size=1,
stride=1,
padding=0,
bias=True,
)
self.submodule1 = Sequential(modules)
self.list1 = [Dropout(drop_prob=0.1), [Softmax(axis=100)]]
self.tuple1 = (
Dropout(drop_prob=0.1),
(Softmax(axis=100), Dropout(drop_prob=0.2)),
)
self.dict1 = {"Dropout": Dropout(drop_prob=0.1)}
self.fc1 = Linear(512, 1024)
def test_syncbn2d_grad():
nr_chan = 8
data_shape = (3, nr_chan, 16, 16)
syncbn = SyncBatchNorm(8, track_running_stats=False)
bn = BatchNorm2d(8, track_running_stats=False)
for i in range(4):
if i == 2:
syncbn.training = False
bn.training = False
inp = Tensor(np.random.normal(loc=2.3, size=data_shape).astype(np.float32))
diff = Tensor(np.random.normal(size=data_shape).astype(np.float32))
with GradManager().attach(inp) as gm:
oup = syncbn(inp)
gm.backward(oup, diff)
grad = inp.grad
inp.grad = None
with GradManager().attach(inp) as gm:
oup_expect = bn(inp)
gm.backward(oup_expect, diff)
grad_expect = inp.grad
inp.grad = None
_assert_allclose(oup.numpy(), oup_expect.numpy())
_assert_allclose(grad.numpy(), grad_expect.numpy())
def test_bn_no_track_stat3():
nchannel = 3
m = BatchNorm2d(nchannel, track_running_stats=False)
m.track_running_stats = True
data = np.random.random((6, nchannel, 2, 2)).astype("float32")
with pytest.raises(Exception):
m(data)
def test_frozen_bn_no_affine(is_training):
nchannel = 3
m = BatchNorm2d(nchannel, freeze=True, affine=False)
if is_training:
m.train()
else:
m.eval()
data = megengine.tensor(np.random.random((6, nchannel, 2, 2)).astype("float32"))
m(data).numpy()
def __init__(self):
super().__init__()
self.conv1 = Conv2d(3, 128, 3, stride=2, bias=False)
self.conv2 = Conv2d(3, 128, 3, padding=1, bias=False)
self.conv3 = Conv2d(3, 128, 3, dilation=2, bias=False)
self.bn1 = BatchNorm2d(128)
self.bn2 = BatchNorm1d(128)
self.dropout = Dropout(drop_prob=0.1)
self.softmax = Softmax(axis=100)
self.pooling = MaxPool2d(kernel_size=2, padding=0)
self.submodule1 = Sequential(Dropout(drop_prob=0.1), Softmax(axis=100),)
self.fc1 = Linear(512, 1024)
def test_bn_no_track_stat():
nchannel = 3
m = BatchNorm2d(nchannel, track_running_stats=False)
gm = ad.GradManager().attach(m.parameters())
optim = optimizer.SGD(m.parameters(), lr=1.0)
optim.clear_grad()
data = np.random.random((6, nchannel, 2, 2)).astype("float32")
with gm:
loss = m(data).sum()
gm.backward(loss)
optim.step()
def test_batchnorm_change_batchsize():
data_shape = (2, 3, 8, 8)
real_shape = (4, 3, 8, 8)
data = np.random.random(data_shape).astype(np.float32)
d = np.random.random(real_shape).astype(np.float32)
bn = BatchNorm2d(3)
f = trace(bn)
f(data)
y1 = f(d)
y0 = bn(tensor(d))
assertTensorClose(y0.numpy(), y1.numpy())
def test_batchnorm2d():
nr_chan = 8
data_shape = (3, nr_chan, 16, 16)
momentum = 0.9
bn = BatchNorm2d(nr_chan, momentum=momentum)
running_mean = np.zeros((1, nr_chan, 1, 1), dtype=np.float32)
running_var = np.ones((1, nr_chan, 1, 1), dtype=np.float32)
data = tensor()
for i in range(3):
xv = np.random.normal(loc=2.3, size=data_shape).astype(np.float32)
xv_transposed = np.transpose(xv, [0, 2, 3, 1]).reshape(
(data_shape[0] * data_shape[2] * data_shape[3], nr_chan))
mean = np.mean(xv_transposed, axis=0).reshape(1, nr_chan, 1, 1)
var_biased = np.var(xv_transposed, axis=0).reshape((1, nr_chan, 1, 1))
sd = np.sqrt(var_biased + bn.eps)
var_unbiased = np.var(xv_transposed, axis=0, ddof=1).reshape(
(1, nr_chan, 1, 1))
running_mean = running_mean * momentum + mean * (1 - momentum)
running_var = running_var * momentum + var_unbiased * (1 - momentum)
data.set_value(xv)
yv = bn(data)
yv_expect = (xv - mean) / sd
assertTensorClose(yv_expect, yv.numpy(), max_err=5e-6)
assertTensorClose(running_mean, bn.running_mean.numpy(), max_err=5e-6)
assertTensorClose(running_var, bn.running_var.numpy(), max_err=5e-6)
# test set 'training' flag to False
mean_backup = bn.running_mean.numpy()
var_backup = bn.running_var.numpy()
bn.training = False
xv = np.random.normal(loc=2.3, size=data_shape).astype(np.float32)
data.set_value(xv)
yv1 = bn(data)
yv2 = bn(data)
assertTensorClose(yv1.numpy(), yv2.numpy(), max_err=0)
assertTensorClose(mean_backup, bn.running_mean.numpy(), max_err=0)
assertTensorClose(var_backup, bn.running_var.numpy(), max_err=0)
yv_expect = (xv - running_mean) / np.sqrt(running_var + bn.eps)
assertTensorClose(yv_expect, yv1.numpy(), max_err=5e-6)
def test_batchnorm2d_no_stats():
nr_chan = 8
data_shape = (3, nr_chan, 16, 16)
bn = BatchNorm2d(8, track_running_stats=False)
for i in range(4):
if i == 2:
bn.training = False
xv = np.random.normal(loc=2.3, size=data_shape).astype(np.float32)
xv_transposed = np.transpose(xv, [0, 2, 3, 1]).reshape(
(data_shape[0] * data_shape[2] * data_shape[3], nr_chan))
mean = np.mean(xv_transposed, axis=0).reshape(1, nr_chan, 1, 1)
var = np.var(xv_transposed, axis=0).reshape((1, nr_chan, 1, 1))
sd = np.sqrt(var + bn.eps)
yv = bn(Tensor(xv))
yv_expect = (xv - mean) / sd
_assert_allclose(yv.numpy(), yv_expect)
def test_batchnorm_empty_tensor(dim, is_symbolic):
if dim == 1:
m = BatchNorm1d(4, affine=True)
inp = mge.tensor(np.random.randn(0, 4, 0).astype("float32"))
elif dim == 2:
m = BatchNorm2d(4, affine=True)
inp = mge.tensor(np.random.randn(0, 4, 0, 0).astype("float32"))
else:
raise NotImplementedError
m.train()
def fn(inp):
return m(inp)
if is_symbolic is not None:
fn = jit.trace(symbolic=is_symbolic)(fn)
for _ in range(3):
out = fn(inp)
np.testing.assert_equal(out.numpy(), inp)
if is_symbolic is None:
break
def test_bn_no_track_stat2():
nchannel = 3
m = BatchNorm2d(nchannel) # Init with track_running_stat = True
m.track_running_stats = False
# m.running_var and m.running_mean created during init time
saved_var = m.running_var.numpy()
assert saved_var is not None
saved_mean = m.running_mean.numpy()
assert saved_mean is not None
gm = ad.GradManager().attach(m.parameters())
optim = optimizer.SGD(m.parameters(), lr=1.0)
optim.clear_grad()
data = np.random.random((6, nchannel, 2, 2)).astype("float32")
with gm:
loss = m(data).sum()
gm.backward(loss)
optim.step()
np.testing.assert_equal(m.running_var.numpy(), saved_var)
np.testing.assert_equal(m.running_mean.numpy(), saved_mean)
def test_frozen_bn():
nchannel = 3
m = BatchNorm2d(nchannel, freeze=True)
saved_var = m.running_var.numpy()
saved_mean = m.running_mean.numpy()
saved_wt = m.weight.numpy()
saved_bias = m.bias.numpy()
gm = ad.GradManager().attach(m.parameters())
optim = optimizer.SGD(m.parameters(), lr=1.0)
optim.clear_grad()
data = np.random.random((6, nchannel, 2, 2)).astype("float32")
with gm:
loss = m(data).mean()
gm.backward(loss)
optim.step()
np.testing.assert_equal(m.running_var.numpy(), saved_var)
np.testing.assert_equal(m.running_mean.numpy(), saved_mean)
np.testing.assert_equal(m.weight.numpy(), saved_wt)
np.testing.assert_equal(m.bias.numpy(), saved_bias)
np.testing.assert_almost_equal(loss.numpy(), data.mean(), 5)
def __init__(self):
super().__init__()
self.bn = BatchNorm2d(4)
def __init__(self):
super().__init__()
self.bn = BatchNorm2d(4)
self.test_bool_key = {True: 1, False: 0}
def __init__(self):
super().__init__()
self.i = self.InnerModule()
self.bn = BatchNorm2d(4)
self.param = Parameter(np.ones(1, dtype=np.float32))
self.buff = Tensor(np.ones(1, dtype=np.float32))
def test_frozen_bn_no_affine():
nchannel = 3
m = BatchNorm2d(nchannel, freeze=True, affine=False)
data = megengine.Tensor(
np.random.random((6, nchannel, 2, 2)).astype("float32"))
m(data).numpy()