|
import matplotlib |
|
matplotlib.use('Agg') |
|
|
|
import functools |
|
|
|
import gradio as gr |
|
import matplotlib.pyplot as plt |
|
import seaborn as sns |
|
import pandas as pd |
|
|
|
|
|
|
|
BENCHMARK_DATA = { |
|
"Greedy Search": { |
|
"DistilGPT2": { |
|
"T4": [336.22, 3976.23, 115.84], |
|
"3090": [158.38, 1835.82, 46.56], |
|
"A100": [371.49, 4073.84, 60.94], |
|
}, |
|
"GPT2": { |
|
"T4": [607.31, 7140.23, 185.12], |
|
"3090": [297.03, 3308.31, 76.68], |
|
"A100": [691.75, 7323.60, 110.72], |
|
}, |
|
"OPT-1.3B": { |
|
"T4": [1303.41, 15939.07, 1488.15], |
|
"3090": [428.33, 7259.43, 468.37], |
|
"A100": [1125.00, 16713.63, 384.52], |
|
}, |
|
"GPTJ-6B": { |
|
"T4": [0, 0, 0], |
|
"3090": [0, 0, 0], |
|
"A100": [2664.28, 32783.09, 1440.06], |
|
}, |
|
"T5 Small": { |
|
"T4": [99.88, 1527.73, 18.78], |
|
"3090": [55.09, 665.70, 9.25], |
|
"A100": [124.91, 1642.07, 13.72], |
|
}, |
|
"T5 Base": { |
|
"T4": [416.56, 6095.05, 106.12], |
|
"3090": [223.00, 2503.28, 46.67], |
|
"A100": [550.76, 6504.11, 64.57], |
|
}, |
|
"T5 Large": { |
|
"T4": [645.05, 9587.67, 225.17], |
|
"3090": [377.74, 4216.41, 97.92], |
|
"A100": [944.17, 10572.43, 116.52], |
|
}, |
|
"T5 3B": { |
|
"T4": [1493.61, 13629.80, 1494.80], |
|
"3090": [694.75, 6316.79, 489.33], |
|
"A100": [1801.68, 16707.71, 411.93], |
|
}, |
|
}, |
|
"Sample": { |
|
"DistilGPT2": { |
|
"T4": [], |
|
"3090": [], |
|
"A100": [], |
|
}, |
|
"GPT2": { |
|
"T4": [], |
|
"3090": [], |
|
"A100": [], |
|
}, |
|
"OPT-1.3B": { |
|
"T4": [], |
|
"3090": [], |
|
"A100": [], |
|
}, |
|
"GPTJ-6B": { |
|
"T4": [], |
|
"3090": [], |
|
"A100": [], |
|
}, |
|
"T5 Small": { |
|
"T4": [], |
|
"3090": [], |
|
"A100": [], |
|
}, |
|
"T5 Base": { |
|
"T4": [], |
|
"3090": [], |
|
"A100": [], |
|
}, |
|
"T5 Large": { |
|
"T4": [], |
|
"3090": [], |
|
"A100": [], |
|
}, |
|
"T5 3B": { |
|
"T4": [], |
|
"3090": [], |
|
"A100": [], |
|
}, |
|
}, |
|
"Beam Search": { |
|
"DistilGPT2": { |
|
"T4": [], |
|
"3090": [], |
|
"A100": [], |
|
}, |
|
"GPT2": { |
|
"T4": [], |
|
"3090": [], |
|
"A100": [], |
|
}, |
|
"OPT-1.3B": { |
|
"T4": [], |
|
"3090": [], |
|
"A100": [], |
|
}, |
|
"GPTJ-6B": { |
|
"T4": [], |
|
"3090": [], |
|
"A100": [], |
|
}, |
|
"T5 Small": { |
|
"T4": [], |
|
"3090": [], |
|
"A100": [], |
|
}, |
|
"T5 Base": { |
|
"T4": [], |
|
"3090": [], |
|
"A100": [], |
|
}, |
|
"T5 Large": { |
|
"T4": [], |
|
"3090": [], |
|
"A100": [], |
|
}, |
|
"T5 3B": { |
|
"T4": [], |
|
"3090": [], |
|
"A100": [], |
|
}, |
|
}, |
|
} |
|
|
|
|
|
def get_plot(model_name, plot_eager, generate_type): |
|
df = pd.DataFrame(BENCHMARK_DATA[generate_type][model_name]) |
|
df["framework"] = ["PyTorch", "TF (Eager Execition)", "TF (XLA)"] |
|
df = pd.melt(df, id_vars=["framework"], value_vars=["T4", "3090", "A100"]) |
|
if plot_eager == "No": |
|
df = df[df["framework"] != "TF (Eager Execition)"] |
|
|
|
g = sns.catplot( |
|
data=df, kind="bar", |
|
x="variable", y="value", hue="framework", |
|
ci="sd", palette="dark", alpha=.6, height=6 |
|
) |
|
g.despine(left=True) |
|
g.set_axis_labels("GPU", "Generation time (ms)") |
|
g.legend.set_title("Framework") |
|
return plt.gcf() |
|
|
|
demo = gr.Blocks() |
|
|
|
with demo: |
|
gr.Markdown( |
|
""" |
|
# TensorFlow XLA Text Generation Benchmark |
|
Pick a tab for the type of generation (or other information), and then select a model from the dropdown menu. |
|
You can also ommit results from TensorFlow Eager Execution, if you wish to better compare the performance of |
|
PyTorch to TensorFlow with XLA. |
|
""" |
|
) |
|
with gr.Tabs(): |
|
with gr.TabItem("Greedy Search"): |
|
gr.Markdown( |
|
""" |
|
### Greedy Search benchmark parameters |
|
- `max_new_tokens = 64`; |
|
- `pad_to_multiple_of = 64` for Tensorflow XLA models. Others do not pad (input prompts between 2 and 33 tokens). |
|
""" |
|
) |
|
with gr.Row(): |
|
model_selector = gr.Dropdown( |
|
choices=["DistilGPT2", "GPT2", "OPT-1.3B", "GPTJ-6B", "T5 Small", "T5 Base", "T5 Large", "T5 3B"], |
|
value="T5 Small", |
|
label="Model", |
|
interactive=True, |
|
) |
|
eager_enabler = gr.Radio( |
|
["Yes", "No"], |
|
value="Yes", |
|
label="Plot TF Eager Execution?", |
|
interactive=True |
|
) |
|
plot_fn = functools.partial(get_plot, generate_type="Greedy Search") |
|
plot = gr.Plot(value=plot_fn("T5 Small", "Yes")) |
|
model_selector.change(fn=plot_fn, inputs=[model_selector, eager_enabler], outputs=plot) |
|
with gr.TabItem("Sample"): |
|
gr.Button("New Tiger") |
|
with gr.TabItem("Beam Search"): |
|
gr.Button("New Tiger") |
|
with gr.TabItem("Benchmark Information"): |
|
gr.Dataframe( |
|
headers=["Parameter", "Value"], |
|
value=[ |
|
["Transformers Version", "4.22.dev0"], |
|
["TensorFlow Version", "2.9.1"], |
|
["Pytorch Version", "1.11.0"], |
|
["OS", "22.04 LTS (3090) / Debian 10 (other GPUs)"], |
|
["CUDA", "11.6 (3090) / 11.3 (others GPUs)"], |
|
["Number of Runs", "100 (the first run was discarded to ignore compilation time)"], |
|
["Is there code to reproduce?", "Yes -- https://gist.github.com/gante/f0017e3f13ac11b0c02e4e4db351f52f"], |
|
], |
|
) |
|
|
|
demo.launch() |
|
|
