def generate_fake_quantization_test(
per_layer: bool = True,
scheme: str = "symmetric",
device_type: str = "cpu",
):
@flow.global_function()
def fake_quantization():
with flow.scope.placement(device_type, "0:0"):
x = flow.get_variable(
name="x1",
shape=(2, 3, 4),
dtype=flow.float,
initializer=flow.random_uniform_initializer(-10, 10),
)
return flow.quantization.fake_quantization(
x,
*flow.quantization.min_max_observer(
x,
per_layer_quantization=per_layer,
quantization_scheme=scheme,
),
quantization_scheme=scheme,
)
convert_to_onnx_and_check(fake_quantization, opset=10 if per_layer else 13)
def generate_min_max_observer_test(
out_pos: int,
per_layer: bool,
formula: str,
scheme: str,
device_type: str = "cpu",
opset: int = 10,
):
@flow.global_function()
def min_max_observer():
with flow.scope.placement(device_type, "0:0"):
x = flow.get_variable(
name="x1",
shape=(2, 3, 4),
dtype=flow.float,
initializer=flow.random_uniform_initializer(-10, 10),
)
return flow.quantization.min_max_observer(
x,
per_layer_quantization=per_layer,
quantization_formula=formula,
quantization_scheme=scheme,
)[out_pos]
convert_to_onnx_and_check(min_max_observer, opset=opset)
def generate_fake_quantization_test_moving_average(
scheme: str = "symmetric", device_type: str = "cpu",
):
@flow.global_function()
def fake_quantization_moving_average():
with flow.scope.placement(device_type, "0:0"):
x = flow.get_variable(
name="x1",
shape=(2, 3, 4),
dtype=flow.float,
initializer=flow.random_uniform_initializer(-10, 10),
)
return flow.quantization.fake_quantization(
x,
*flow.quantization.moving_average_min_max_observer(
x, quantization_scheme=scheme,
),
quantization_scheme=scheme,
)
set_moving_max_min_value()
convert_to_onnx_and_check(
fake_quantization_moving_average, opset=10, explicit_init=False
)
def test_conv2d_k2d1_valid(test_case):
@flow.global_function(func_config)
def conv2d_k3s1_valid(x=flow.FixedTensorDef((2, 4, 3, 5))):
return flow.layers.conv2d(
x, 6, kernel_size=3, strides=1, padding="VALID", **initer_args
)
convert_to_onnx_and_check(conv2d_k3s1_valid)
def test_conv2d_s2_same(test_case):
@flow.global_function(func_config)
def conv2d_s2_same(x=flow.FixedTensorDef((2, 4, 3, 5))):
return flow.layers.conv2d(
x, 6, kernel_size=3, strides=2, padding="SAME", **initer_args
)
convert_to_onnx_and_check(conv2d_s2_same)
def test_resnet50(test_case):
@flow.global_function()
def InferenceNet(images: tp.Numpy.Placeholder((1, 3, 224, 224))):
logits = resnet50(images)
predictions = flow.nn.softmax(logits)
return predictions
convert_to_onnx_and_check(InferenceNet)
def test_mobilenetv2(test_case):
func_config = flow.FunctionConfig()
func_config.default_data_type(flow.float)
@flow.global_function(func_config)
def mobilenetv2(x=flow.FixedTensorDef((1, 224, 224, 3))):
return Mobilenet(x)
convert_to_onnx_and_check(mobilenetv2)
def test_alexnet(test_case):
func_config = flow.FunctionConfig()
func_config.default_data_type(flow.float)
@flow.global_function(func_config)
def alexnet_eval_job(x=flow.FixedTensorDef((1, 227, 227, 3))):
return alexnet(x, None, False)
convert_to_onnx_and_check(alexnet_eval_job)
def test_conv2d_s2_same(test_case):
@flow.global_function()
def conv2d_s2_same(x: tp.Numpy.Placeholder((2, 4, 3, 5))):
return flow.layers.conv2d(x,
6,
kernel_size=3,
strides=2,
padding="SAME",
**initer_args)
convert_to_onnx_and_check(conv2d_s2_same)
def test_conv2d_k2d1_valid(test_case):
@flow.global_function()
def conv2d_k3s1_valid(x: tp.Numpy.Placeholder((2, 4, 3, 5))):
return flow.layers.conv2d(x,
6,
kernel_size=3,
strides=1,
padding="VALID",
**initer_args)
convert_to_onnx_and_check(conv2d_k3s1_valid)
def test_bias_add_nhwc(test_case):
@flow.global_function()
def bias_add_nhwc(x: tp.Numpy.Placeholder((3, 4, 2, 5))):
y = flow.get_variable(
name="y",
shape=(5,),
dtype=flow.float,
initializer=flow.random_uniform_initializer(),
)
return flow.nn.bias_add(x, y, "NHWC")
convert_to_onnx_and_check(bias_add_nhwc)
def test_resnet50(test_case):
func_config = flow.FunctionConfig()
func_config.default_data_type(flow.float)
@flow.global_function(func_config)
def InferenceNet(images=flow.FixedTensorDef((1, 3, 224, 224))):
logits = resnet50(images)
predictions = flow.nn.softmax(logits)
return predictions
convert_to_onnx_and_check(InferenceNet)
def test_bias_add_nhwc(test_case):
@flow.global_function(func_config)
def bias_add_nhwc(x=flow.FixedTensorDef((3, 4, 2, 5))):
y = flow.get_variable(
name="y",
shape=(5, ),
dtype=flow.float,
initializer=flow.random_uniform_initializer(),
)
return flow.nn.bias_add(x, y, "NHWC")
convert_to_onnx_and_check(bias_add_nhwc)
def test_conv2d_k3s1_nchw_same_g2(test_case):
@flow.global_function()
def conv2d(x: tp.Numpy.Placeholder((2, 4, 5, 3))):
return flow.layers.conv2d(x,
6,
kernel_size=3,
strides=1,
groups=2,
padding="SAME",
data_format="NCHW",
**initer_args)
convert_to_onnx_and_check(conv2d)
def test_conv2d_k3s1_nhwc_same_d2(test_case):
@flow.global_function()
def conv2d(x: tp.Numpy.Placeholder((2, 7, 5, 4))):
return flow.layers.conv2d(x,
6,
kernel_size=3,
strides=1,
dilation_rate=2,
padding="SAME",
data_format="NHWC",
**initer_args)
convert_to_onnx_and_check(conv2d)
def test_max_pooling_2d_k3s1_valid_nhwc(test_case):
@flow.global_function()
def max_pooling_2d_k3s1_valid_nhwc(x: tp.Numpy.Placeholder((2, 3, 5, 4))):
x += flow.get_variable(
name="v1",
shape=(1, 1),
dtype=flow.float,
initializer=flow.zeros_initializer(),
)
return flow.nn.max_pool2d(
x, ksize=3, strides=1, padding="VALID", data_format="NHWC"
)
convert_to_onnx_and_check(max_pooling_2d_k3s1_valid_nhwc)
def test_max_pooling_2d_k3s1_valid_nchw(test_case):
@flow.global_function(func_config)
def max_pooling_2d_k3s1_valid_nchw(x=flow.FixedTensorDef((2, 3, 5, 4))):
x += flow.get_variable(
name="v1",
shape=(1, 1),
dtype=flow.float,
initializer=flow.zeros_initializer(),
)
return flow.nn.max_pool2d(
x, ksize=3, strides=1, padding="VALID", data_format="NCHW"
)
convert_to_onnx_and_check(max_pooling_2d_k3s1_valid_nchw)
def test_max_pooling_2d_k2s2_same_nhwc(test_case):
@flow.global_function(func_config)
def max_pooling_2d_k2s2_same_nhwc(x=flow.FixedTensorDef((2, 3, 5, 4))):
x += flow.get_variable(
name="v1",
shape=(1, 1),
dtype=flow.float,
initializer=flow.zeros_initializer(),
)
return flow.nn.max_pool2d(
x, ksize=2, strides=2, padding="SAME", data_format="NHWC"
)
convert_to_onnx_and_check(max_pooling_2d_k2s2_same_nhwc)
def test_max_pooling_2d_k2s2_same_nchw(test_case):
@flow.global_function()
def max_pooling_2d_k2s2_same_nchw(x: tp.Numpy.Placeholder((2, 3, 5, 4))):
x += flow.get_variable(
name="v1",
shape=(1, 1),
dtype=flow.float,
initializer=flow.zeros_initializer(),
)
return flow.nn.max_pool2d(
x, ksize=2, strides=2, padding="SAME", data_format="NCHW"
)
convert_to_onnx_and_check(max_pooling_2d_k2s2_same_nchw)
def test_add_many(test_case):
@flow.global_function(func_config)
def add_many():
variables = []
for i in range(50):
variables.append(
flow.get_variable(
name=str(i),
shape=(2, 3),
dtype=flow.float,
initializer=flow.random_uniform_initializer(),
))
return flow.math.add_n(variables)
convert_to_onnx_and_check(add_many)
def test_concat_axis1(test_case):
@flow.global_function(func_config)
def concat():
variables = []
for i in range(4):
variables.append(
flow.get_variable(
name=str(i),
shape=(2, 3),
dtype=flow.float,
initializer=flow.random_uniform_initializer(),
))
return flow.concat(variables, axis=1)
convert_to_onnx_and_check(concat)
def test_conv2d_k3s1_nchw_same_g2(test_case):
@flow.global_function(func_config)
def conv2d(x=flow.FixedTensorDef((2, 4, 5, 3))):
return flow.layers.conv2d(
x,
6,
kernel_size=3,
strides=1,
groups=2,
padding="SAME",
data_format="NCHW",
**initer_args
)
convert_to_onnx_and_check(conv2d)
def test_conv2d_k3s1_nhwc_same_d2(test_case):
@flow.global_function(func_config)
def conv2d(x=flow.FixedTensorDef((2, 7, 5, 4))):
return flow.layers.conv2d(
x,
6,
kernel_size=3,
strides=1,
dilation_rate=2,
padding="SAME",
data_format="NHWC",
**initer_args
)
convert_to_onnx_and_check(conv2d)
def generate_unary_op_test(flow_op,
*args,
opset=None,
min_val=-10,
max_val=10,
**kwargs):
@flow.global_function()
def job1():
x = flow.get_variable(
name="x1",
shape=(2, 3, 4),
dtype=flow.float,
initializer=flow.random_uniform_initializer(min_val, max_val),
)
return flow_op(x, *args, **kwargs)
convert_to_onnx_and_check(job1, opset=opset)
def test_add_2(test_case):
@flow.global_function()
def add_2():
x = flow.get_variable(
name="x",
shape=(2, 3),
dtype=flow.float,
initializer=flow.random_uniform_initializer(),
)
y = flow.get_variable(
name="y",
shape=(2, 3),
dtype=flow.float,
initializer=flow.random_uniform_initializer(),
)
return flow.math.add_n([x, y])
convert_to_onnx_and_check(add_2)
def test_batch_matmul(test_case):
@flow.global_function(func_config)
def matmul():
a = flow.get_variable(
name="a",
shape=(4, 2, 3),
dtype=flow.float,
initializer=flow.random_uniform_initializer(),
)
b = flow.get_variable(
name="b",
shape=(4, 3, 4),
dtype=flow.float,
initializer=flow.random_uniform_initializer(),
)
return flow.matmul(a, b)
convert_to_onnx_and_check(matmul)
def test_matmul_ta_tb(test_case):
@flow.global_function(func_config)
def matmul():
a = flow.get_variable(
name="a",
shape=(3, 2),
dtype=flow.float,
initializer=flow.random_uniform_initializer(),
)
b = flow.get_variable(
name="b",
shape=(4, 3),
dtype=flow.float,
initializer=flow.random_uniform_initializer(),
)
return flow.matmul(a, b, transpose_a=True, transpose_b=True)
convert_to_onnx_and_check(matmul)
def generate_binary_op_test(flow_op, *args, opset=None, **kwargs):
@flow.global_function(func_config)
def job1():
x = flow.get_variable(
name="x1",
shape=(2, 3, 4),
dtype=flow.float,
initializer=flow.random_uniform_initializer(-10, 10),
)
y = flow.get_variable(
name="y1",
shape=(1, 3, 1),
dtype=flow.float,
initializer=flow.random_uniform_initializer(-10, 10),
)
return flow_op(x, y, *args, **kwargs)
convert_to_onnx_and_check(job1, opset=opset)
def test_gather_nd(test_case):
@flow.global_function(func_config)
def gather_nd():
x = flow.get_variable(
name="x",
shape=(2, 3, 4),
dtype=flow.float,
initializer=flow.random_uniform_initializer(),
)
y = flow.get_variable(
name="y",
shape=(2, 3),
dtype=flow.int64,
initializer=flow.random_uniform_initializer(0, 1, flow.int64),
)
return flow.gather_nd(x, y)
convert_to_onnx_and_check(gather_nd, opset=11)
def test_large_array(test_case):
@flow.global_function(func_config)
def add_with_large_array():
large_shape = (256 * 1024 * 1024 + 1, )
x = flow.get_variable(
name="x",
shape=large_shape,
dtype=flow.float,
initializer=flow.random_uniform_initializer(),
)
y = flow.get_variable(
name="y",
shape=large_shape,
dtype=flow.float,
initializer=flow.random_uniform_initializer(),
)
return flow.math.add_n([x, y])
# ONNX Runtime optimizers doesn't support external data
convert_to_onnx_and_check(add_with_large_array,
external_data=True,
ort_optimize=False)