Spaces:
Sleeping
Sleeping
File size: 3,157 Bytes
198d05b 1ef37f4 198d05b 1ef37f4 198d05b 1ef37f4 198d05b 1ef37f4 198d05b 1ef37f4 198d05b 1ef37f4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 |
import re
import pandas as pd
import gradio as gr
from py_markdown_table.markdown_table import markdown_table
from model_sizer.utils import get_sizes, create_empty_model, convert_bytes
def convert_url_to_name(url:str):
"Converts a model URL to its name on the Hub"
results = re.findall(r"huggingface.co\/(.*?)#", url)
if len(results) < 1:
raise ValueError(f"URL {url} is not a valid model URL to the Hugging Face Hub")
return results[0]
def calculate_memory(model_name:str, library:str, options:list):
"Calculates the memory usage for a model"
if library == "auto":
library = None
if "huggingface.co" in model_name:
model_name = convert_url_to_name(model_name)
model = create_empty_model(model_name, library_name=library)
total_size, largest_layer = get_sizes(model)
data = []
title = f"Memory Usage for `{model_name}`"
for dtype in options:
dtype_total_size = total_size
dtype_largest_layer = largest_layer[0]
if dtype == "float16":
dtype_total_size /= 2
dtype_largest_layer /= 2
elif dtype == "int8":
dtype_total_size /= 4
dtype_largest_layer /= 4
elif dtype == "int4":
dtype_total_size /= 8
dtype_largest_layer /= 8
dtype_training_size = convert_bytes(dtype_total_size * 4)
dtype_total_size = convert_bytes(dtype_total_size)
dtype_largest_layer = convert_bytes(dtype_largest_layer)
data.append({
"dtype": dtype,
"Largest Layer": dtype_largest_layer,
"Total Size": dtype_total_size,
"Training using Adam": dtype_training_size
})
return f'## {title}', pd.DataFrame(data)
with gr.Blocks() as demo:
gr.Markdown(
"""# Model Memory Calculator
This tool will help you calculate how much vRAM is needed to train and perform big model inference
on a model hosted on the :hugging_face: Hugging Face Hub. The minimum recommended vRAM needed for a model
is denoted as the size of the "largest layer", and training of a model is roughly 4x its size (for Adam).
Currently this tool supports all models hosted that use `transformers` and `timm`.
To use this tool pass in the URL or model name of the model you want to calculate the memory usage for,
select which framework it originates from ("auto" will try and detect it from the model metadata), and
what precisions you want to use.
"""
)
out_text = gr.Markdown()
out = gr.DataFrame(
headers=["dtype", "Largest Layer", "Total Size", "Training using Adam"],
)
inp = gr.Textbox(label="Model Name or URL")
with gr.Row():
library = gr.Radio(["auto", "transformers", "timm"], label="Library", value="auto")
options = gr.CheckboxGroup(
["float32", "float16", "int8", "int4"],
value="float32"
)
btn = gr.Button("Calculate Memory Usage", scale=0.5)
btn.click(
calculate_memory, inputs=[inp, library, options], outputs=[out_text, out],
)
demo.launch() |