INNER CODE UNIT · Python
speed_test
Ki6an/fastT5 · fastT5/model_testing_tools.py:8
def speed_test(
onnx_model,
torch_model,
beam_range: range = range(1, 10, 1),
seq_length_range: range = range(10, 500, 50),
input_text=None,
):
"""
method prints the time took for onnx and pytorch model to finish a text generation task
args:
input_text (str) : text input for the model.
onnx_model : onnx representation of the t5 model,
torch_model : torch represention of the t5 model,
beam_range (range) : provide a range, which takes starting end and steps (don't start with 0)
sequence_length-range (range) : takes the start, end and steps as a range (start with 10)
return :
onnx_model_latency : numpy array of latency for each beam number and sequence length