def test_reshape_fused():
out = tensor(np.arange(100, 116, dtype=np.int32).reshape(1, 16))
out = out.reshape(tensor(2), 2, tensor(4), 1)
assertTensorClose(
out.numpy(), np.arange(100, 116, dtype=np.int32).reshape(2, 2, 4, 1)
)
def test_many_batch_interpolate():
inp = tensor(np.arange(1, 9, dtype=np.float32).reshape(2, 1, 2, 2))
out = F.interpolate(inp, [4, 4])
out2 = F.interpolate(inp, scale_factor=2.0)
assertTensorClose(out.numpy(), out2.numpy())
def __call__(self, ori_params, new_params, step):
for param in new_params:
grad = param.grad.numpy()
self.s_slots[param] += grad**2
delta = grad / (self.s_slots[param] + self.eps)**0.5
delta *= -(self.lr / (1 + (step - 1) * self.lr_decay))
assertTensorClose(param.numpy(), ori_params[param] + delta)
def test_assign_corner_interpolate():
inp = tensor(np.arange(1, 5, dtype=np.float32).reshape(1, 1, 2, 2))
out = F.interpolate(inp, [4, 4], align_corners=True)
out2 = F.interpolate(inp, scale_factor=2.0, align_corners=True)
assertTensorClose(out.numpy(), out2.numpy())
def test_onehot_low_dimension():
inp = tensor(np.arange(1, 4, dtype=np.int32))
out = F.one_hot(inp)
assertTensorClose(
out.numpy(),
np.eye(4, dtype=np.int32)[np.arange(1, 4, dtype=np.int32)])
def test_eager_equvilence():
eager_net = SimpleNet()
trace_enable_net = copy.deepcopy(eager_net)
trace_disable_net = copy.deepcopy(eager_net)
opt_factory = lambda net: SGD(
net.parameters(requires_grad=True), lr=0.01, momentum=0.01
)
estep = generate_eager_step(eager_net, opt_factory)
te_step = generate_trace_step(trace_enable_net, opt_factory, True)
td_step = generate_trace_step(trace_disable_net, opt_factory, False)
assert_network_equvilence([eager_net, trace_enable_net, trace_disable_net])
# Use hard code number as limit, may increase if needed.
for data, label in itertools.islice(minibatch_generator(), 200):
eloss = estep(data, label)
te_loss = te_step(data, label)
td_loss = td_step(data, label)
assertTensorClose(eloss, te_loss)
assertTensorClose(eloss, td_loss)
assert_network_equvilence(
[eager_net, trace_enable_net, trace_disable_net,]
)
def test_conv_transpose2d():
SH, SW = 3, 1
PH, PW = 2, 0
N, IC, IH, IW = 4, 5, 8, 6
KH, KW = 3, 4
OC = 3
BIAS = True
def getsize(inp, kern, stride):
return (inp - 1) * stride + kern
OH = getsize(IH, KH, SH)
OW = getsize(IW, KW, SW)
inp = np.random.normal(size=(N, IC, IH, IW)).astype(np.float32)
out = np.zeros((N, OC, OH, OW), dtype=np.float32)
weight = np.random.normal(size=(IC, OC, KH, KW)).astype(np.float32)
bias = np.random.normal(size=(1, OC, 1, 1)).astype(np.float32)
# naive calculation use numpy
for n, ic, ih, iw in itertools.product(*map(range, [N, IC, IH, IW])):
oh, ow = ih * SH, iw * SW
out[n, :, oh : oh + KH, ow : ow + KW] += inp[n, ic, ih, iw] * weight[ic]
out = out[:, :, PH : OH - PH, PW : OW - PW]
if BIAS:
out += bias
# megengine conv_transpose2d calculation
conv_transpose2d = ConvTranspose2d(IC, OC, (KH, KW), (SH, SW), (PH, PW), bias=BIAS)
conv_transpose2d.weight = Parameter(weight, dtype=np.float32)
if BIAS:
conv_transpose2d.bias = Parameter(bias, dtype=np.float32)
y = conv_transpose2d(tensor(inp))
assertTensorClose(out, y.numpy(), max_err=2e-6)
def test_abs():
assertTensorClose(
F.abs(tensor([-3.0, -4.0, -5.0])).numpy(),
np.abs(np.array([-3.0, -4.0, -5.0], dtype=np.float32)),
)
assertTensorClose(F.abs(-3.0), np.abs(np.float32(-3.0)))
def test_sgd_momentum_static():
_, data_shape, _, label_shape = get_input()
mlp = MLP()
opt = SGD(mlp.parameters(), lr=0.01, momentum=0.9)
@trace
def f(data, label):
pred = mlp(data)
loss = F.square_loss(pred, label.reshape(-1, 1))
opt.zero_grad()
opt.backward(loss)
slots = TensorDict()
for param in mlp.parameters():
slots[param] = np.zeros(param.shape).astype(np.float32)
for _ in range(3):
f(
np.random.random(data_shape).astype(np.float32),
np.random.randint(0, 10, label_shape).astype(np.int32),
)
orig_params = TensorDict()
grads = TensorDict()
for param in mlp.parameters():
orig_params[param] = np.copy(param.numpy())
grads[param] = np.copy(param.grad.numpy())
opt.step()
for param in mlp.parameters():
slot = slots[param]
orig_param = orig_params[param]
slot *= 0.9
slot -= param.grad.numpy() * 0.01
assertTensorClose(param.numpy(), orig_param + slot)
def test_sgd_simple():
data, data_shape, label, label_shape = get_input()
mlp = MLP()
opt = SGD(mlp.parameters(), lr=0.01, weight_decay=0.1)
for idx in range(3):
data.set_value(np.random.random(data_shape).astype(np.float32))
label.set_value(np.random.randint(0, 10, label_shape))
pred = mlp(data)
loss = F.square_loss(pred, label.reshape(-1, 1))
if idx % 2:
opt.zero_grad()
else:
mlp.zero_grad()
opt.backward(loss)
grads = TensorDict()
orig_params = TensorDict()
for param in mlp.parameters():
grad = F.grad(loss, param, use_virtual_grad=False)
assertTensorClose(grad.numpy(), param.grad.numpy())
grads[param] = np.copy(grad.numpy())
orig_params[param] = np.copy(param.numpy())
opt.step()
for param in mlp.parameters():
assertTensorClose(param.numpy(),
orig_params[param] * 0.999 - grads[param] * 0.01)
def test_pytorch_backward():
class APlusB(torch.nn.Module):
def __init__(self):
super(APlusB, self).__init__()
def forward(self, a, b):
return a + b
a = randomTorch(15, 15)
b = randomTorch(15, 15)
def get_pytorch_backward():
parameter_a = a.clone()
parameter_a.requires_grad = True
c = APlusB()(parameter_a, b)
d = APlusB()(c, b)
e = torch.sum(d)
e.backward()
return parameter_a.grad
def get_mge_backward():
mge_module = PyTorchModule(APlusB())
mge_a = Parameter(a.numpy(), dtype=np.float32)
mge_b = tensor(b.numpy(), dtype=np.float32)
mge_c = mge_module(mge_a, mge_b)
mge_d = mge_module(mge_c, mge_b)
mge_e = mge.functional.sum(mge_d)
return mge.functional.grad(mge_e, mge_a, use_virtual_grad=False)
assertTensorClose(get_pytorch_backward().numpy(), get_mge_backward().numpy())
def test_update_lr():
data, data_shape, label, label_shape = get_input()
mlp = MLP()
opt = SGD(mlp.parameters(), lr=0.01)
pred = mlp(data)
loss = F.square_loss(pred, label.reshape(-1, 1))
opt.zero_grad()
opt.backward(loss)
opt.step()
for group in opt.param_groups:
group["lr"] += 0.02
for _ in range(3):
data.set_value(np.random.random(data_shape).astype(np.float32))
label.set_value(np.random.randint(0, 10, label_shape))
pred = mlp(data)
loss = F.square_loss(pred, label.reshape(-1, 1))
opt.zero_grad()
opt.backward(loss)
for param in mlp.parameters():
grad = F.grad(loss, param, use_virtual_grad=False)
assertTensorClose(grad.numpy(), param.grad.numpy())
orig_params = []
for param in mlp.parameters():
orig_params.append(np.copy(param.numpy()))
opt.step()
for param, orig_param in zip(mlp.parameters(), orig_params):
assertTensorClose(param.numpy(),
orig_param - param.grad.numpy() * 0.03)
def test_compile_multi_times_static():
return # XXX: rewrite or remove this test
with Graph() as cg:
cg.set_option("eager_evaluation", False)
data = Input("data", shape=(2, 28))
label = Input("label", shape=(2, ), dtype=np.int32)
mlp = MLP()
opt = SGD(mlp.parameters(requires_grad=True), lr=0.01)
pred0 = mlp(data)
pred = F.softmax(pred0)
loss = F.square_loss(pred, label.reshape(2, 1))
opt.zero_grad()
grads = opt.backward(loss)
opt.step()
f0 = compile(pred, None)
f1 = compile([pred, loss], grads, copy=True)
data = np.random.random((2, 28)).astype(np.float32)
label = np.random.randint(0, 10, (2, )).astype(np.float32)
out0 = f0(data=data)
out1 = f1(data=data, label=label)
assertTensorClose(out0[0], out1[0])
_ = compile([pred, loss], grads, copy=False)
with pytest.raises(mgb.MegBrainError):
f0(data=data)
def test_load_quantized():
data_shape = (2, 28)
data = tensor(np.random.random(data_shape), dtype="float32")
data = data.astype(mgb.dtype.qint8(0.1))
mlp = MLP()
quantize_qat(mlp)
quantize(mlp)
mlp.dense0.weight = Parameter(
mlp.dense0.weight.astype(mgb.dtype.qint8(0.001)).numpy())
mlp.dense1.weight = Parameter(
mlp.dense1.weight.astype(mgb.dtype.qint8(0.0002)).numpy())
mlp.eval()
pred0 = mlp(data)
with BytesIO() as fout:
mge.save(mlp.state_dict(), fout)
fout.seek(0)
checkpoint = mge.load(fout)
# change mlp weight.
mlp.dense0.weight = Parameter(
mlp.dense0.weight.astype(mgb.dtype.qint8(0.00001)).numpy())
mlp.dense1.weight = Parameter(
mlp.dense1.weight.astype(mgb.dtype.qint8(0.2)).numpy())
mlp.load_state_dict(checkpoint)
pred1 = mlp(data)
assertTensorClose(pred0.astype("float32").numpy(),
pred1.astype("float32").numpy(),
max_err=5e-6)
def test_leaky_relu():
data = np.array([-8, -12, 6, 10]).astype(np.float32)
negative_slope = 0.1
leaky_relu = LeakyReLU(negative_slope)
output = leaky_relu(mge.tensor(data))
np_output = np.maximum(0, data) + negative_slope * np.minimum(0, data)
assertTensorClose(output.numpy(), np_output, max_err=0)
def test_ones():
assertTensorClose(
mge.ones((2, 2), dtype=np.int32).numpy(),
np.ones((2, 2), dtype=np.int32))
assertTensorClose(
mge.ones(mge.tensor([2, 2], dtype=np.int32), dtype=np.int32).numpy(),
np.ones((2, 2), dtype=np.int32),
)
def test_onehot_high_dimension():
arr = np.array(
[[3, 2, 4, 4, 2, 4, 0, 4, 4, 1], [4, 1, 1, 3, 2, 2, 4, 2, 4, 3]],
dtype=np.int32)
inp = tensor(arr)
out = F.one_hot(inp, 10)
assertTensorClose(out.numpy(), np.eye(10, dtype=np.int32)[arr])
def test_clamp():
"""Fix an issue when `lower` or `upper` is 0, it will be recognized as `False` and
`F.clamp` will fall into wrong conditions unexpectedly.
"""
x = np.linspace(-6, 6, dtype="float32")
assertTensorClose(
F.clamp(tensor(x) + 3, 0, 6).numpy(), np.clip(x + 3, 0, 6))
assertTensorClose(
F.clamp(tensor(x) - 3, -6, 0).numpy(), np.clip(x - 3, -6, 0))
def test_local_conv2d():
batch_size = 10
in_channels = 4
out_channels = 8
input_height = 8
input_width = 8
kernel_size = 3
stride = 1
padding = 1
dilation = 1
groups = 1
local_conv2d = LocalConv2d(
in_channels=in_channels,
out_channels=out_channels,
input_height=input_height,
input_width=input_width,
kernel_size=kernel_size,
stride=stride,
padding=padding,
dilation=dilation,
groups=groups,
)
inputs = np.random.normal(
size=(batch_size, in_channels, input_height, input_width)
).astype(np.float32)
output_height = (input_height + padding * 2 - kernel_size) // stride + 1
output_width = (input_width + padding * 2 - kernel_size) // stride + 1
weights = np.random.normal(
size=(
groups,
output_height,
output_width,
in_channels // groups,
kernel_size,
kernel_size,
out_channels // groups,
)
).astype(np.float32)
local_conv2d.weight = Parameter(weights)
outputs = local_conv2d(tensor(inputs))
# naive calculation use numpy
# only test output_height == input_height, output_width == input_width, group == 1
inputs = np.pad(inputs, ((0, 0), (0, 0), (1, 1), (1, 1)))
expected = np.zeros(
(batch_size, out_channels, output_height, output_width), dtype=np.float32,
)
for n, oc, oh, ow in itertools.product(
*map(range, [batch_size, out_channels, output_height, output_width])
):
ih, iw = oh * stride, ow * stride
expected[n, oc, ih, iw] = np.sum(
inputs[n, :, ih : ih + kernel_size, iw : iw + kernel_size]
* weights[0, oh, ow, :, :, :, oc]
)
assertTensorClose(outputs.numpy(), expected, max_err=1e-5)
def test_shape_infer():
@jit.trace(symbolic=True)
def f(x):
a, b = x.shape
return sum(x[i] for i in range(a))
x = np.random.randn(3, 10).astype("float32")
assertTensorClose(f(x), x.sum(0))
x = np.random.randn(4, 10).astype("float32")
assertTensorClose(f(x), x[:3].sum(0))
def __call__(self, ori_params, new_params, step):
for param in new_params:
grad = param.grad.numpy()
if hasattr(self, "momentum"):
self.slots[
param] = grad + self.slots[param] * self.momentum
delta = -self.lr * self.slots[param]
else:
delta = -self.lr * grad
assertTensorClose(param.numpy(), ori_params[param] + delta)
def assert_network_equvilence(nets):
net_state = [net.state_dict() for net in nets]
for state in net_state[1:]:
assert len(net_state[0]) == len(state)
for k, v in net_state[0].items():
for state in net_state[1:]:
assert k in state
assertTensorClose(v, state[k])
def test_broadcast_shapeof():
inp = tensor(np.arange(1, 17, dtype=np.int32).reshape(4, 4))
out = tensor(np.arange(100, 104, dtype=np.int32).reshape(1, 4))
out = out.broadcast(inp.shapeof())
tmp = np.array([[100, 101, 102, 103]], dtype=np.int32)
out2 = np.repeat(tmp, 4, axis=0)
assertTensorClose(out.numpy(), out2)
def test_linear_interpolate():
inp = tensor(np.arange(1, 3, dtype=np.float32).reshape(1, 1, 2))
out = F.interpolate(inp, scale_factor=2.0, mode="LINEAR")
out2 = F.interpolate(inp, 4, mode="LINEAR")
assertTensorClose(out.numpy(),
np.array([[[1.0, 1.25, 1.75, 2.0]]], dtype=np.float32))
assertTensorClose(out2.numpy(),
np.array([[[1.0, 1.25, 1.75, 2.0]]], dtype=np.float32))
def test_set_value():
v0 = np.random.random((2, 3)).astype(np.float32)
param = Parameter(v0)
v1 = np.random.random((2, 3)).astype(np.float32)
param.set_value(v1)
assertTensorClose(param.numpy(), v1, max_err=5e-6)
v2 = np.random.random((3, 3)).astype(np.float32)
# TODO: add this
# with pytest.raises(ValueError):
# param.set_value(v2)
assertTensorClose(param.numpy(), v1, max_err=5e-6)
def __call__(self, ori_params, new_params, step):
for param in new_params:
grad = param.grad.numpy()
self.s_slots[param] = self.s_slots[
param] * self.rho + grad**2 * (1 - self.rho)
delta = (grad * ((self.a_slots[param] + self.eps)**0.5) /
(self.s_slots[param] + self.eps)**0.5)
self.a_slots[param] = self.a_slots[
param] * self.rho + delta**2 * (1 - self.rho)
delta *= -self.lr
assertTensorClose(param.numpy(), ori_params[param] + delta)
def train_test(backend):
model_path = "../examples/cifar10/resnet_example/checkpoint/pretrained_model_82.mge"
# Change the reference number if the change is from numerical rounding-off
# FIXME! Need to use different number depending on CPU/GPU
if backend == "megengine-dynamic":
os.environ["MGE_DISABLE_TRACE"] = "true"
loss_ref = np.array([3.4709125, 12.46342]).astype(np.float32)
else:
loss_ref = np.array([3.4709125, 12.463419]).astype(np.float32)
import megengine
from megengine.functional.debug_param import set_conv_execution_strategy
from megengine.test import assertTensorClose
from megengine.core import Graph
sys.path.append(
os.path.join(os.path.dirname(__file__), "..", "..", "..", "examples"))
from cifar10.resnet_example.main import Example as resnet18_config
from cifar10.resnet_example.main import train_one_iter_mge
mge_root = os.path.dirname(megengine.__file__)
model_path = os.path.join(mge_root, model_path)
set_conv_execution_strategy("HEURISTIC_REPRODUCIBLE")
run_case = resnet18_config(backend=backend, mode="train")
run_case.init_net()
run_case.load_model(model_path)
max_err = 0.0
loss = []
np.random.seed(0)
inputs = np.random.rand(run_case.train_batch_size, 3, 32, 32)
targets = np.random.randint(10, size=(run_case.train_batch_size, ))
run_case.set_optimizer(0.0)
opt = run_case.net_context["optimizer"]
for lr in (1.0, 1.0):
run_case.set_optimizer(lr)
opt.zero_grad()
loss_batch, _ = train_one_iter_mge(inputs, targets, config=run_case)
opt.step()
loss.append(loss_batch.numpy()[0])
try:
assertTensorClose(np.array(loss).astype(np.float32),
loss_ref,
max_err=1e-5)
except:
print("calculated loss:", loss)
print("expect:", loss_ref)
sys.exit(1)
def run(use_trace, symbolic):
a = tensor(np.array([1926.0817], dtype=np.float32))
net = Sigmoid()
func_run = run_saved_context
if use_trace:
func_run = trace(run_saved_context, symbolic=symbolic)
s = func_run(a, net=net)
s2 = F.sigmoid(a)
assertTensorClose(s.numpy(), s2.numpy())
assertTensorClose(
F.grad(s, a, use_virtual_grad=False).numpy(),
F.grad(s2, a, use_virtual_grad=False).numpy(),
)
def check_value():
for param in mlp.parameters():
grad = param.grad.numpy()
orig_param = orig_params[param]
m = m_slots[param]
v = v_slots[param]
m *= beta0
m += (1 - beta0) * grad
v *= beta1
v += (1 - beta1) * grad * grad
update = (m / (1 - beta0**step_size)) / (
np.sqrt(v / (1 - beta1**step_size)) + eps)
assertTensorClose(param.numpy(), orig_param - 0.01 * update)
def __call__(self, ori_params, new_params, step):
for param in new_params:
grad = param.grad.numpy()
m = self.m_slots[param]
v = self.v_slots[param]
m *= self.betas[0]
m += (1 - self.betas[0]) * grad
v *= self.betas[1]
v += (1 - self.betas[1]) * grad * grad
delta = (m / (1 - self.betas[0]**step)) / (
np.sqrt(v / (1 - self.betas[1]**step)) + self.eps)
assertTensorClose(param.numpy(),
ori_params[param] - self.lr * delta)