def test_slice():
net = SubtensorOpr()
mge_result = dump_mge_model(net, net.data, tmp_file)
_test_convert_result(net.data,
tmp_file,
mge_result,
max_error,
nhwc=False,
nhwc2=False)
def test_bn():
net = BnOpr("bn1d")
net.eval()
mge_result = dump_mge_model(net, net.data1, tmp_file, True)
_test_convert_result(net.data1,
tmp_file,
mge_result,
max_error,
disable_nhwc=True)
def test_xornet():
if megengine.__version__ < "1.1.0":
return
net = XORNet("tflite")
net.eval()
mge_result = dump_mge_model(net, net.data, tmp_file, True)
_test_convert_result(net.data,
tmp_file,
mge_result,
max_error,
disable_nhwc=True)
def test_model(model):
data = np.ones((1, 3, 224, 224)).astype(np.float32)
if megengine.__version__ < "1.1.0":
commit_id = "dc2f2cfb228a135747d083517b98aea56e7aab92"
else:
commit_id = None
net = megengine.hub.load("megengine/models",
model,
use_cache=False,
commit=commit_id,
pretrained=True)
mge_result = dump_mge_model(net, data, tmp_file)
_test_convert_result(data, tmp_file, mge_result, 1e-4)
def test_elemwise_broadcast(mode):
net = ElemwiseOpr(mode)
mge_result = dump_mge_model(net,
np.array([2.0]).astype("float32"), tmp_file)
_test_convert_result(
np.array([2.0]).astype("float32"), tmp_file, mge_result, max_error)
def test_typecvt():
net = TypeCvtOpr()
mge_result = dump_mge_model(net, net.data, tmp_file)
_test_convert_result(net.data, tmp_file, mge_result, max_error)
def test_leakyrelu_model():
if megengine.__version__ < "1.1.0":
return
net = XORNet()
mge_result = dump_mge_model(net, net.data, tmp_file, False)
_test_convert_result(net.data, tmp_file, mge_result, max_error)
def test_convbn():
net = ConvBn2dOpr()
mge_result = dump_mge_model(net, net.data, tmp_file)
_test_convert_result(net.data, tmp_file, mge_result, max_error)
def test_active(mode):
net = ActiveOpr(mode, fused=False)
mge_result = dump_mge_model(net, net.data, tmp_file)
_test_convert_result(net.data, tmp_file, mge_result, max_error)
def test_reduce(mode):
net = ReduceOpr(mode=mode)
mge_result = dump_mge_model(net, net.data, tmp_file)
_test_convert_result(net.data, tmp_file, mge_result, max_error)
def test_reshape():
net = ReshapeOpr(fix_batch=True)
mge_result = dump_mge_model(net, net.data, tmp_file)
_test_convert_result(net.data, tmp_file, mge_result, max_error, nhwc=False)
def test_xornet():
if mge.__version__ < "1.1.0":
return
net = XORNet()
mge_result = dump_mge_model(net, net.data, tmp_file, True)
_test_convert_result(net.data, tmp_file, mge_result, max_error)
def test_broadcast():
net = BroadcastOpr()
mge_result = dump_mge_model(net, net.data, tmp_file)
_test_convert_result(net.data, tmp_file, mge_result, max_error, min_version=8)
def test_subtensor():
net = SubtensorOpr()
mge_result = dump_mge_model(net, net.data, tmp_file)
_test_convert_result(net.data, tmp_file, mge_result, max_error)
def test_pool(mode):
if mode == "avg":
return
net = PoolOpr(mode)
mge_result = dump_mge_model(net, net.data, tmp_file)
_test_convert_result(net.data, tmp_file, mge_result, max_error)
def test_batchnorm(mode):
net = BnOpr(mode)
net.eval()
data = net.data1 if mode == "bn1d" else net.data2
mge_result = dump_mge_model(net, data, tmp_file)
_test_convert_result(data, tmp_file, mge_result, max_error)
def test_softmax():
net = SoftmaxOpr()
mge_result = dump_mge_model(net, net.data, tmp_file)
_test_convert_result(net.data, tmp_file, mge_result, max_error)
def test_pooling(mode):
if megengine.__version__ > "0.6.0" and mode == "avg":
return
net = PoolOpr(mode)
mge_result = dump_mge_model(net, net.data, tmp_file)
_test_convert_result(net.data, tmp_file, mge_result, max_error)
def test_linear():
net = LinearOpr()
mge_result = dump_mge_model(net, net.data, tmp_file)
_test_convert_result(net.data, tmp_file, mge_result, max_error)
def test_transopse():
net = TransposeOpr()
mge_result = dump_mge_model(net, net.data, tmp_file)
_test_convert_result(net.data, tmp_file, mge_result, max_error)
def test_conv2d(mode):
net = ConvOpr(mode)
mge_result = dump_mge_model(net, net.data, tmp_file)
_test_convert_result(net.data, tmp_file, mge_result, max_error)
def test_squeeze():
net = SqueezeOpr()
mge_result = dump_mge_model(net, net.data, tmp_file)
_test_convert_result(net.data, tmp_file, mge_result, max_error)
def test_elemwise(mode):
net = ElemwiseOpr(mode)
mge_result = dump_mge_model(net, net.data, tmp_file)
_test_convert_result(net.data, tmp_file, mge_result, max_error)
from test.utils import ConvOpr, dump_mge_model
import megengine as mge
import numpy as np
from megengine.core.tensor import dtype
from megengine.quantization.quantize import quantize_qat
from megengine.traced_module import trace_module
if __name__ == "__main__":
net = ConvOpr("normal")
traced_module = trace_module(net, mge.tensor(net.data))
mge.save(traced_module, "float_model.tm")
dump_mge_model(net, net.data, "float_model")
qat_net = quantize_qat(net)
inp_dtype = dtype.qint8(16.0 / 128)
data = mge.tensor(np.random.random((1, 3, 224, 224))) * 16
data = data.astype(inp_dtype)
inp = mge.tensor(dtype.convert_from_qint8(data.numpy()))
inp.qparams.scale = mge.tensor(dtype.get_scale(inp_dtype))
inp.qparams.dtype_meta = dtype._builtin_quant_dtypes["qint8"]
qat_module = trace_module(qat_net, inp)
mge.save(qat_module, "qat_model.tm")
