saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-22 19:23:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add parity test for LayerNormalization (#8622)

Browse source
This commit is contained in:
Tianlei Wu 2021-08-05 10:11:19 -07:00 committed by GitHub
parent dda9f53bed
commit 24b14c650b
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 437 additions and 231 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
onnxruntime/python/tools/transformers/fusion_layernorm.py
View file

@ -41,15 +41,13 @@ class FusionLayerNormalization(Fusion):
if len(children) == 0 or len(children) > 2:
return
parent = self.model.get_parent(node, 0, output_name_to_node)
if parent is None:
return
if children[0].op_type != 'Sub' or self.model.get_parent(children[0], 0, output_name_to_node) != parent:
root_input = node.input[0]
if children[0].op_type != 'Sub' or children[0].input[0] != root_input:
return
if len(children) == 2:
if children[1].op_type != 'Sub' or self.model.get_parent(children[1], 0, output_name_to_node) != parent:
if children[1].op_type != 'Sub' or children[1].input[0] != root_input:
return
div_node = None

156
onnxruntime/test/python/transformers/parity_utilities.py Normal file
View file

@ -0,0 +1,156 @@
# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# -------------------------------------------------------------------------
import torch
from torch import nn
import numpy
import os
import sys
def create_inputs(batch_size=1, sequence_length=1, hidden_size=768, float16=False, device=torch.device('cuda')):
float_type = torch.float16 if float16 else torch.float32
input = torch.normal(mean=0.0, std=1.0, size=(batch_size, sequence_length, hidden_size)).to(float_type).to(device)
return input
def export_onnx(model, onnx_model_path, float16, hidden_size, device):
from pathlib import Path
Path(onnx_model_path).parent.mkdir(parents=True, exist_ok=True)
input_hidden_states = create_inputs(hidden_size=hidden_size, float16=float16, device=device)
with torch.no_grad():
outputs = model(input_hidden_states)
dynamic_axes = {'input': {0: 'batch_size', 1: 'seq_len'}, "output": {0: 'batch_size', 1: 'seq_len'}}
torch.onnx.export(model,
args=(input_hidden_states),
f=onnx_model_path,
input_names=['input'],
output_names=["output"],
dynamic_axes=dynamic_axes,
example_outputs=outputs,
opset_version=11,
do_constant_folding=True)
print("exported:", onnx_model_path)
def optimize_onnx(input_onnx_path, optimized_onnx_path, expected_op=None):
# Try import optimizer from source directory so that we need not build and install package after making change.
source_dir = os.path.join(os.path.dirname(__file__), '..', '..', '..', 'python', 'tools', 'transformers')
if (os.path.exists(source_dir) and source_dir not in sys.path):
sys.path.append(source_dir)
from optimizer import optimize_model
else:
from onnxruntime.transformers.optimizer import optimize_model
onnx_model = optimize_model(input_onnx_path, model_type='gpt2', opt_level=0)
onnx_model.save_model_to_file(optimized_onnx_path)
if expected_op is not None:
assert len(onnx_model.get_nodes_by_op_type(expected_op)) == 1, \
f"Expected {expected_op} node not found in the optimized model {optimized_onnx_path}"
def diff_outputs(torch_outputs, ort_outputs, index):
""" Returns the maximum difference between PyTorch and OnnxRuntime outputs.
"""
expected_outputs = torch_outputs[index].cpu().numpy()
diff = numpy.abs(expected_outputs - ort_outputs[index])
return numpy.amax(diff)
def compare_outputs(torch_outputs, ort_outputs, atol=1e-06, verbose=True):
"""Compare outputs from PyTorch and OnnxRuntime
Args:
torch_outputs (Tuple[Torch.Tensor]): PyTorch model output
ort_outputs (List[numpy.ndarray]): OnnxRuntime output
atol (float, optional): Absolute tollerance. Defaults to 1e-06.
verbose (bool, optional): Print more information. Defaults to True.
Returns:
is_all_close(bool): whether all elements are close.
max_abs_diff(float): maximum absolute difference.
"""
same = numpy.asarray([
numpy.allclose(ort_outputs[i], torch_outputs[i].cpu().numpy(), atol=atol, rtol=0)
for i in range(len(ort_outputs))
])
max_abs_diff = [diff_outputs(torch_outputs, ort_outputs, i) for i in range(len(ort_outputs))]
is_all_close = same.all()
if (not is_all_close) and verbose:
for i in numpy.where(numpy.logical_not(same))[0]:
diff = numpy.fabs(ort_outputs[i] - torch_outputs[i].cpu().numpy())
idx = numpy.unravel_index(diff.argmax(), diff.shape)
print(
f'Output {i}, diff={diff[idx]:.9f} index={idx} ort={ort_outputs[i][idx]:.9f} torch={float(torch_outputs[i][idx]):.9f}'
)
return is_all_close, max(max_abs_diff)
def create_ort_session(onnx_model_path, use_gpu=True):
from onnxruntime import SessionOptions, InferenceSession, GraphOptimizationLevel, __version__ as onnxruntime_version
sess_options = SessionOptions()
sess_options.graph_optimization_level = GraphOptimizationLevel.ORT_DISABLE_ALL
sess_options.intra_op_num_threads = 2
sess_options.log_severity_level = 2
execution_providers = ['CPUExecutionProvider'] if not use_gpu else ['CUDAExecutionProvider', 'CPUExecutionProvider']
return InferenceSession(onnx_model_path, sess_options, providers=execution_providers)
def onnxruntime_inference(ort_session, input):
ort_inputs = {'input': numpy.ascontiguousarray(input.cpu().numpy())}
ort_outputs = ort_session.run(None, ort_inputs)
return ort_outputs
def run_parity(model,
onnx_model_path,
batch_size,
hidden_size,
sequence_length,
float16,
device,
optimized,
test_cases=100,
verbose=False,
tolerance=None):
passed_cases = 0
max_diffs = []
printed = False # print only one sample
ort_session = create_ort_session(onnx_model_path, device.type == 'cuda')
for i in range(test_cases):
input_hidden_states = create_inputs(batch_size, sequence_length, hidden_size, float16, device)
with torch.no_grad():
torch_outputs = model(input_hidden_states)
ort_outputs = onnxruntime_inference(ort_session, input_hidden_states)
if tolerance is None:
tolerance = 1e-03 if float16 else 1e-05
is_all_close, max_diff = compare_outputs(torch_outputs, ort_outputs, atol=tolerance, verbose=verbose)
max_diffs.append(max_diff)
if is_all_close:
passed_cases += 1
elif verbose and not printed:
printed = True
numpy.set_printoptions(precision=10, floatmode='fixed')
torch.set_printoptions(precision=10)
print("input", input_hidden_states)
print("torch_outputs", torch_outputs)
print("ort_outputs", ort_outputs)
max_diff = max(max_diffs)
diff_count = len([i for i in max_diffs if i > 0])
success_flag = "[FAILED]" if passed_cases < test_cases else "[OK]"
print(f"{success_flag} Passed_cases={passed_cases}/{test_cases}; Max_diff={max_diff}; Diff_count={diff_count}")
return test_cases - passed_cases

149
onnxruntime/test/python/transformers/test_parity_gelu.py
View file

@ -26,9 +26,9 @@ For comparison, CPU has MaxDiff=4.77E-07 for each formula.
import unittest
import torch
from torch import nn
import numpy
import math
import os
from parity_utilities import *
class Gelu(nn.Module):
@ -68,141 +68,14 @@ class Gelu(nn.Module):
return (fp32_gelu, )
def create_inputs(batch_size=1, sequence_length=1, hidden_size=768, float16=False, device=torch.device('cuda')):
float_type = torch.float16 if float16 else torch.float32
input = torch.normal(mean=0.0, std=1.0, size=(batch_size, sequence_length, hidden_size)).to(float_type).to(device)
return input
def get_output_names():
outputs = ["output"]
return outputs
def export_onnx(model, onnx_model_path, float16, hidden_size, device):
from pathlib import Path
Path(onnx_model_path).parent.mkdir(parents=True, exist_ok=True)
input_hidden_states = create_inputs(hidden_size=hidden_size, float16=float16, device=device)
with torch.no_grad():
outputs = model(input_hidden_states)
dynamic_axes = {'input': {0: 'batch_size', 1: 'seq_len'}, "output": {0: 'batch_size', 1: 'seq_len'}}
torch.onnx.export(model,
args=(input_hidden_states),
f=onnx_model_path,
input_names=['input'],
output_names=["output"],
dynamic_axes=dynamic_axes,
example_outputs=outputs,
opset_version=11,
do_constant_folding=True)
print("exported:", onnx_model_path)
def optimize_onnx(input_onnx_path, optimized_onnx_path, expected_gelu_op_type='Gelu'):
from onnxruntime.transformers.optimizer import optimize_model
onnx_model = optimize_model(input_onnx_path, model_type='gpt2', opt_level=0)
assert len(onnx_model.get_nodes_by_op_type(
expected_gelu_op_type)) == 1, f"Expected {expected_gelu_op_type} node not found in the optimized model"
onnx_model.save_model_to_file(optimized_onnx_path)
def diff_outputs(torch_outputs, ort_outputs, index):
""" Returns the maximum difference between PyTorch and OnnxRuntime outputs.
"""
expected_outputs = torch_outputs[index].cpu().numpy()
diff = numpy.abs(expected_outputs - ort_outputs[index])
return numpy.amax(diff)
def compare_outputs(torch_outputs, ort_outputs, atol=1e-06, verbose=True):
""" Returns True if torch and ORT outputs are close for given thresholds, and False otherwise.
"""
same = numpy.asarray([
numpy.allclose(ort_outputs[i], torch_outputs[i].cpu().numpy(), atol=atol, rtol=0)
for i in range(len(ort_outputs))
])
max_abs_diff = [diff_outputs(torch_outputs, ort_outputs, i) for i in range(len(ort_outputs))]
is_all_close = same.all()
if is_all_close:
for i in numpy.where(numpy.logical_not(same))[0]:
diff = numpy.fabs(ort_outputs[i] - torch_outputs[i].cpu().numpy())
idx = numpy.unravel_index(diff.argmax(), diff.shape)
print(
f'Output {i}, diff={diff[idx]:.9f} index={idx} ort={ort_outputs[i][idx]:.9f} torch={float(torch_outputs[i][idx]):.9f}'
)
return is_all_close, max(max_abs_diff)
def create_ort_session(onnx_model_path, use_gpu=True):
from onnxruntime import SessionOptions, InferenceSession, GraphOptimizationLevel, __version__ as onnxruntime_version
sess_options = SessionOptions()
sess_options.graph_optimization_level = GraphOptimizationLevel.ORT_DISABLE_ALL
sess_options.intra_op_num_threads = 2
sess_options.log_severity_level = 2
execution_providers = ['CPUExecutionProvider'] if not use_gpu else ['CUDAExecutionProvider', 'CPUExecutionProvider']
return InferenceSession(onnx_model_path, sess_options, providers=execution_providers)
def onnxruntime_inference(ort_session, input):
ort_inputs = {'input': numpy.ascontiguousarray(input.cpu().numpy())}
ort_outputs = ort_session.run(None, ort_inputs)
return ort_outputs
def run_parity(model,
onnx_model_path,
batch_size,
hidden_size,
sequence_length,
float16,
device,
optimized,
test_cases=100,
verbose=False):
print(
f"optimized={optimized}, onnx_model_path={onnx_model_path}, batch_size={batch_size}, hidden_size={hidden_size}, sequence_length={sequence_length}, float16={float16}, device={device}"
)
passed_cases = 0
max_diffs = []
printed = False # print only one sample
ort_session = create_ort_session(onnx_model_path, device.type == 'cuda')
for i in range(test_cases):
input_hidden_states = create_inputs(batch_size, sequence_length, hidden_size, float16, device)
with torch.no_grad():
torch_outputs = model(input_hidden_states)
ort_outputs = onnxruntime_inference(ort_session, input_hidden_states)
tolerance = 1e-04 if float16 else 1e-06
is_all_close, max_diff = compare_outputs(torch_outputs, ort_outputs, atol=tolerance)
max_diffs.append(max_diff)
if is_all_close:
passed_cases += 1
elif verbose and not printed:
printed = True
numpy.set_printoptions(precision=10, floatmode='fixed')
torch.set_printoptions(precision=10)
print("input", input_hidden_states)
print("torch_outputs", torch_outputs)
print("ort_outputs", ort_outputs)
max_diff = max(max_diffs)
diff_count = len([i for i in max_diffs if i > 0])
success_flag = "[FAILED]" if passed_cases < test_cases else "[OK]"
print(f"{success_flag} Passed_cases={passed_cases}/{test_cases}; Max_diff={max_diff}; Diff_count={diff_count}")
return test_cases - passed_cases
def run(batch_size, float16, optimized, hidden_size, device, test_cases, formula=0, sequence_length=2):
test_name = f"batch_size={batch_size}, float16={float16}, optimized={optimized}, hidden_size={hidden_size}, formula={formula}"
print(f"\nTesting ONNX parity: {test_name}")
test_name = f"device={device}, float16={float16}, optimized={optimized}, batch_size={batch_size}, sequence_length={sequence_length}, hidden_size={hidden_size}, formula={formula}"
print(f"\nTesting: {test_name}")
model = Gelu(formula=formula)
model.eval()
@ -216,7 +89,7 @@ def run(batch_size, float16, optimized, hidden_size, device, test_cases, formula
if optimized:
optimized_onnx_path = './temp/gelu_{}_opt_{}.onnx'.format(formula, "fp16" if float16 else "fp32")
optimize_onnx(onnx_model_path, optimized_onnx_path, expected_gelu_op_type=Gelu.get_fused_op(formula))
optimize_onnx(onnx_model_path, optimized_onnx_path, Gelu.get_fused_op(formula))
onnx_path = optimized_onnx_path
else:
onnx_path = onnx_model_path
@ -238,14 +111,16 @@ class TestGeluParity(unittest.TestCase):
self.test_cases = 100 # Number of test cases per test run
self.sequence_length = 2
self.hidden_size = 768
self.formula_to_test = [0, 1, 3, 4, 5] # formula 2 cannot pass precision test.
self.formula_to_test = [0, 1, 2, 3, 4, 5]
self.formula_must_pass = [0, 1, 3, 4, 5] # formula 2 cannot pass precision test.
def run_test(self, batch_size, float16, optimized, hidden_size, device, formula):
def run_test(self, batch_size, float16, optimized, hidden_size, device, formula, enable_assert=True):
if float16 and device.type == 'cpu': # CPU does not support FP16
return
num_failure, test_name = run(batch_size, float16, optimized, hidden_size, device, self.test_cases, formula,
self.sequence_length)
self.assertTrue(num_failure == 0, test_name)
if enable_assert:
self.assertTrue(num_failure == 0, "Failed: " + test_name)
def run_one(self, optimized, device, hidden_size=768, formula=0):
for batch_size in [4]:
@ -254,14 +129,16 @@ class TestGeluParity(unittest.TestCase):
optimized=optimized,
hidden_size=hidden_size,
device=device,
formula=formula)
formula=formula,
enable_assert=formula in self.formula_must_pass)
self.run_test(batch_size,
float16=True,
optimized=optimized,
hidden_size=hidden_size,
device=device,
formula=formula)
formula=formula,
enable_assert=formula in self.formula_must_pass)
def test_cpu(self):
cpu = torch.device('cpu')

126
onnxruntime/test/python/transformers/test_parity_huggingface_gpt_attention.py
View file

@ -20,6 +20,7 @@ import onnx
import numpy
import os
from transformers.modeling_utils import Conv1D
from parity_utilities import diff_outputs, create_ort_session, compare_outputs
DEBUG_OUTPUTS = ["qk", "norm_qk", "softmax", "attn_weights"]
@ -284,51 +285,6 @@ def optimize_onnx(input_onnx_path, optimized_onnx_path, num_heads, debug):
onnx_model.save_model_to_file(optimized_onnx_path)
def diff_outputs(torch_outputs, ort_outputs, index, relative=False):
""" Returns the maximum difference between PyTorch and OnnxRuntime outputs.
"""
expected_outputs = torch_outputs[index].cpu().numpy()
diff = numpy.abs(expected_outputs - ort_outputs[index])
if relative:
return numpy.amax(diff / (numpy.abs(expected_outputs) + 1e-6))
else:
return numpy.amax(diff)
def compare_outputs(torch_outputs, ort_outputs, rtol=1e-03, atol=1e-03, verbose=False):
""" Returns True if torch and ORT outputs are close for given thresholds, and False otherwise.
"""
is_all_close = True
max_abs_diff = []
for i in range(len(ort_outputs)):
is_close = numpy.allclose(ort_outputs[i], torch_outputs[i].cpu().numpy(), rtol=rtol, atol=atol)
if not is_close:
is_all_close = False
if verbose:
print(f'output {i} ({ort_outputs[i].name}) are close: {is_close}')
max_abs_diff.append(diff_outputs(torch_outputs, ort_outputs, i))
if (not is_all_close) or (verbose and max(max_abs_diff) > 0):
messages = ["max_abs_diff per output:"]
output_names = ["attn_output", "present"] + DEBUG_OUTPUTS
for i, diff in enumerate(max_abs_diff):
messages.append(f"{output_names[i]}={diff},")
print(" ".join(messages))
return is_all_close, max(max_abs_diff)
def create_ort_session(onnx_model_path, use_gpu=True):
from onnxruntime import SessionOptions, InferenceSession, GraphOptimizationLevel, __version__ as onnxruntime_version
sess_options = SessionOptions()
sess_options.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_ALL
sess_options.intra_op_num_threads = 2
sess_options.log_severity_level = 2
execution_providers = ['CPUExecutionProvider'
] if not use_gpu else ['CUDAExecutionProvider', 'CPUExecutionProvider']
return InferenceSession(onnx_model_path, sess_options, providers=execution_providers)
def onnxruntime_inference(ort_session, input_hidden_states, attention_mask, past):
ort_inputs = {
'past': numpy.ascontiguousarray(past.cpu().numpy()),
@ -370,7 +326,7 @@ def verify_attention(model,
ort_outputs = onnxruntime_inference(ort_session, input_hidden_states, attention_mask, layer_past)
tolerance = 1e-03 if float16 else 1e-05
is_all_close, max_diff = compare_outputs(torch_outputs, ort_outputs, rtol=tolerance, atol=tolerance)
is_all_close, max_diff = compare_outputs(torch_outputs, ort_outputs, atol=tolerance, verbose=True)
max_diffs.append(max_diff)
if is_all_close:
passed_cases += 1
@ -410,52 +366,22 @@ def run(batch_size, float16, optimized, hidden_size, num_attention_heads, device
past_sequence_length = 0
padding_length = 0
num_failure = 0
num_failure += verify_attention(model,
onnx_path,
batch_size,
hidden_size,
num_attention_heads,
sequence_length,
past_sequence_length,
float16,
device,
padding_length,
optimized,
test_cases)
num_failure += verify_attention(model, onnx_path, batch_size, hidden_size, num_attention_heads, sequence_length,
past_sequence_length, float16, device, padding_length, optimized, test_cases)
# Test Case: with past state and padding last 2 words
sequence_length = 3
past_sequence_length = 5
padding_length = 2
num_failure += verify_attention(model,
onnx_path,
batch_size,
hidden_size,
num_attention_heads,
sequence_length,
past_sequence_length,
float16,
device,
padding_length,
optimized,
test_cases)
num_failure += verify_attention(model, onnx_path, batch_size, hidden_size, num_attention_heads, sequence_length,
past_sequence_length, float16, device, padding_length, optimized, test_cases)
# Test Case: random mask one word
sequence_length = 1
past_sequence_length = 128
padding_length = -1
num_failure += verify_attention(model,
onnx_path,
batch_size,
hidden_size,
num_attention_heads,
sequence_length,
past_sequence_length,
float16,
device,
padding_length,
optimized,
test_cases)
num_failure += verify_attention(model, onnx_path, batch_size, hidden_size, num_attention_heads, sequence_length,
past_sequence_length, float16, device, padding_length, optimized, test_cases)
# clean up onnx file
os.remove(onnx_model_path)
@ -467,24 +393,45 @@ def run(batch_size, float16, optimized, hidden_size, num_attention_heads, device
class TestGptAttentionHuggingfaceParity(unittest.TestCase):
def setUp(self):
self.optimized = True # Change it to False if you want to test parity of non optimized ONNX
self.optimized = True # Change it to False if you want to test parity of non optimized ONNX
self.test_cases = 10 # Number of test cases per test run
def run_test(self, batch_size, float16, optimized, hidden_size, num_attention_heads, device):
if float16 and device.type=='cpu': # CPU does not support FP16
if float16 and device.type == 'cpu': # CPU does not support FP16
return
num_failure, test_name = run(batch_size, float16, optimized, hidden_size, num_attention_heads, device, self.test_cases)
num_failure, test_name = run(batch_size, float16, optimized, hidden_size, num_attention_heads, device,
self.test_cases)
self.assertTrue(num_failure == 0, test_name)
def run_small(self, optimized, device):
for batch_size in [64]:
self.run_test(batch_size, float16=False, optimized=optimized, hidden_size=768, num_attention_heads=12, device=device)
self.run_test(batch_size, float16=True, optimized=optimized, hidden_size=768, num_attention_heads=12, device=device)
self.run_test(batch_size,
float16=False,
optimized=optimized,
hidden_size=768,
num_attention_heads=12,
device=device)
self.run_test(batch_size,
float16=True,
optimized=optimized,
hidden_size=768,
num_attention_heads=12,
device=device)
def run_large(self, optimized, device):
for batch_size in [2]:
self.run_test(batch_size, float16=False, optimized=optimized, hidden_size=4096, num_attention_heads=32, device=device)
self.run_test(batch_size, float16=True, optimized=optimized, hidden_size=4096, num_attention_heads=32, device=device)
self.run_test(batch_size,
float16=False,
optimized=optimized,
hidden_size=4096,
num_attention_heads=32,
device=device)
self.run_test(batch_size,
float16=True,
optimized=optimized,
hidden_size=4096,
num_attention_heads=32,
device=device)
def test_cpu(self):
cpu = torch.device('cpu')
@ -507,5 +454,6 @@ class TestGptAttentionHuggingfaceParity(unittest.TestCase):
gpu = torch.device('cuda')
self.run_large(self.optimized, gpu)
if __name__ == '__main__':
unittest.main()

227
onnxruntime/test/python/transformers/test_parity_layernorm.py Normal file
View file

@ -0,0 +1,227 @@
# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
import unittest
import torch
from torch import nn
import os
from parity_utilities import *
class LayerNorm(nn.Module):
def __init__(self, hidden_size, epsilon, cast_fp16=True):
super().__init__()
self.layer_norm = nn.LayerNorm(hidden_size, eps=epsilon)
# initialize weights with random value
self.layer_norm.bias.data.normal_(mean=0.0, std=0.1)
self.layer_norm.weight.data.normal_(mean=0.0, std=0.5)
self.cast_fp16 = cast_fp16
@staticmethod
def get_fused_op():
return "LayerNormalization"
def forward(self, x):
if self.cast_fp16 and x.dtype == torch.float16:
y = self.layer_norm(x.to(torch.float32)).to(torch.float16)
return (y, )
y = self.layer_norm(x)
return (y, )
def optimize_fp16_onnx_with_cast(input_onnx_path, optimized_onnx_path, epsilon):
from onnxruntime.transformers.onnx_model import OnnxModel
import onnx
m = onnx.load(input_onnx_path)
onnx_model = OnnxModel(m)
nodes_to_remove = onnx_model.nodes()
nodes_to_add = [
onnx.helper.make_node("Cast", ["input"], ["fp32_input"], "cast_input", to=1),
onnx.helper.make_node("Cast", ["layer_norm.weight"], ["fp32_layer_norm.weight"], "cast_weight", to=1),
onnx.helper.make_node("Cast", ["layer_norm.bias"], ["fp32_layer_norm.bias"], "cast_bias", to=1),
onnx.helper.make_node("LayerNormalization", ["fp32_input", "fp32_layer_norm.weight", "fp32_layer_norm.bias"],
["fp32_output"],
"layer_norm",
epsilon=epsilon), # use fp32 epsilon
onnx.helper.make_node("Cast", ["fp32_output"], ["output"], "cast_output", to=10)
]
onnx_model.remove_nodes(nodes_to_remove)
onnx_model.add_nodes(nodes_to_add)
onnx_model.prune_graph()
onnx_model.save_model_to_file(optimized_onnx_path)
def optimize_fp16_onnx_no_cast(input_onnx_path, optimized_onnx_path, epsilon):
from onnxruntime.transformers.onnx_model import OnnxModel
import onnx
m = onnx.load(input_onnx_path)
onnx_model = OnnxModel(m)
nodes_to_remove = onnx_model.nodes()
node_to_add = onnx.helper.make_node("LayerNormalization", ["input", "layer_norm.weight", "layer_norm.bias"],
["output"],
"layer_norm",
epsilon=epsilon)
onnx_model.remove_nodes(nodes_to_remove)
onnx_model.add_node(node_to_add)
onnx_model.prune_graph()
onnx_model.save_model_to_file(optimized_onnx_path)
def get_output_names():
outputs = ["output"]
return outputs
def run(batch_size,
float16,
optimized,
hidden_size,
device,
test_cases,
sequence_length=2,
epsilon=0.00001,
cast_fp16=True,
cast_onnx_only=False,
verbose=False):
test_name = f"device={device}, float16={float16}, optimized={optimized}, batch_size={batch_size}, sequence_length={sequence_length}, hidden_size={hidden_size}, epsilon={epsilon}, cast_fp16={cast_fp16}, cast_onnx_only={cast_onnx_only}"
print(f"\nTesting: {test_name}")
model = LayerNorm(hidden_size, epsilon, cast_fp16)
model.eval()
model.to(device)
if float16 and not cast_fp16:
model.half()
# Do not re-use onnx file from previous test since weights of model are random.
onnx_model_path = './temp/layer_norm_{}.onnx'.format("fp16" if float16 else "fp32")
export_onnx(model, onnx_model_path, float16, hidden_size, device)
if optimized:
optimized_onnx_path = './temp/layer_norm_{}_opt.onnx'.format("fp16" if float16 else "fp32")
if (not float16) or cast_fp16:
optimize_onnx(onnx_model_path, optimized_onnx_path,
expected_op=LayerNorm.get_fused_op())
else:
if cast_onnx_only:
optimize_fp16_onnx_with_cast(onnx_model_path, optimized_onnx_path, epsilon=epsilon)
else:
optimize_fp16_onnx_no_cast(onnx_model_path, optimized_onnx_path, epsilon=epsilon)
onnx_path = optimized_onnx_path
else:
onnx_path = onnx_model_path
num_failure = run_parity(model,
onnx_path,
batch_size,
hidden_size,
sequence_length,
float16,
device,
optimized,
test_cases,
verbose=verbose)
# clean up onnx file
os.remove(onnx_model_path)
if optimized:
os.remove(onnx_path)
return num_failure, test_name
class TestLayerNormParity(unittest.TestCase):
def setUp(self):
self.optimized = True # Change it to False if you want to test parity of non optimized ONNX
self.test_cases = 100 # Number of test cases per test run
self.sequence_length = 2
self.hidden_size = 768
self.epsilon = 0.00001
self.verbose = False
def run_test(self,
batch_size,
float16,
optimized,
hidden_size,
device,
cast_fp16=True,
cast_onnx_only=False,
enable_assert=True):
if float16 and device.type == 'cpu': # CPU does not support FP16
return
num_failure, test_name = run(batch_size,
float16,
optimized,
hidden_size,
device,
self.test_cases,
self.sequence_length,
self.epsilon,
cast_fp16,
cast_onnx_only,
verbose=self.verbose)
if enable_assert:
self.assertTrue(num_failure == 0, "Failed: " + test_name)
def run_one(self, optimized, device, hidden_size=768):
for batch_size in [4]:
self.run_test(batch_size, float16=False, optimized=optimized, hidden_size=hidden_size, device=device)
self.run_test(
batch_size,
float16=True,
optimized=optimized,
hidden_size=hidden_size,
device=device,
cast_fp16=True,
cast_onnx_only=False,
enable_assert=False # This setting has small chance to exceed tollerance threshold 0.001
)
self.run_test(
batch_size,
float16=True,
optimized=optimized,
hidden_size=hidden_size,
device=device,
cast_fp16=False,
cast_onnx_only=False,
enable_assert=False # This setting cannot pass tollerance threshold
)
self.run_test(
batch_size,
float16=True,
optimized=optimized,
hidden_size=hidden_size,
device=device,
cast_fp16=False,
cast_onnx_only=True,
enable_assert=False # This setting cannot pass tollerance threshold
)
def test_cpu(self):
cpu = torch.device('cpu')
self.run_one(self.optimized, cpu, hidden_size=self.hidden_size)
def test_cuda(self):
if not torch.cuda.is_available():
import pytest
pytest.skip('test requires GPU and torch+cuda')
else:
gpu = torch.device('cuda')
self.run_one(self.optimized, gpu, hidden_size=self.hidden_size)
if __name__ == '__main__':
unittest.main()