mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-12 17:57:38 +00:00
Update stable diffusion benchmark script (#14759)
Update stable diffusion benchmark script: (1) Test GPU memory usage (2) Change diffusers version to 0.13, and add support of PyTorch 2.0 including compile (3) Add support of xformers (4) Output result to CSV file Example to run PyTorch 2.0 with torch.compile: ``` pip3 install numpy --pre torch --force-reinstall --extra-index-url https://download.pytorch.org/whl/nightly/cu117 export TRITON_PTXAS_PATH=/usr/local/cuda-11.7/bin/ptxas python benchmark.py -e torch -v 1.5 -c 5 -n 1 -b 1 --enable_torch_compile ```
This commit is contained in:
Tianlei Wu
GitHub
parent
1b7f65437e
commit
262e46e8ce
2 changed files with 250 additions and 27 deletions
|
|
@ -4,7 +4,10 @@
|
|||
# --------------------------------------------------------------------------
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import os
|
||||
import statistics
|
||||
import sys
|
||||
import time
|
||||
|
||||
SD_MODELS = {
|
||||
|
|
@ -31,6 +34,94 @@ def example_prompts():
|
|||
return prompts
|
||||
|
||||
|
||||
def measure_gpu_memory(func, start_memory=None):
|
||||
class MemoryMonitor:
|
||||
def __init__(self, keep_measuring=True):
|
||||
self.keep_measuring = keep_measuring
|
||||
|
||||
def measure_gpu_usage(self):
|
||||
from py3nvml.py3nvml import (
|
||||
NVMLError,
|
||||
nvmlDeviceGetCount,
|
||||
nvmlDeviceGetHandleByIndex,
|
||||
nvmlDeviceGetMemoryInfo,
|
||||
nvmlDeviceGetName,
|
||||
nvmlInit,
|
||||
nvmlShutdown,
|
||||
)
|
||||
|
||||
max_gpu_usage = []
|
||||
gpu_name = []
|
||||
try:
|
||||
nvmlInit()
|
||||
device_count = nvmlDeviceGetCount()
|
||||
if not isinstance(device_count, int):
|
||||
print(f"nvmlDeviceGetCount result is not integer: {device_count}")
|
||||
return None
|
||||
|
||||
max_gpu_usage = [0 for i in range(device_count)]
|
||||
gpu_name = [nvmlDeviceGetName(nvmlDeviceGetHandleByIndex(i)) for i in range(device_count)]
|
||||
while True:
|
||||
for i in range(device_count):
|
||||
info = nvmlDeviceGetMemoryInfo(nvmlDeviceGetHandleByIndex(i))
|
||||
if isinstance(info, str):
|
||||
print(f"nvmlDeviceGetMemoryInfo returns str: {info}")
|
||||
return None
|
||||
max_gpu_usage[i] = max(max_gpu_usage[i], info.used / 1024**2)
|
||||
time.sleep(0.002) # 2ms
|
||||
if not self.keep_measuring:
|
||||
break
|
||||
nvmlShutdown()
|
||||
return [
|
||||
{
|
||||
"device_id": i,
|
||||
"name": gpu_name[i],
|
||||
"max_used_MB": max_gpu_usage[i],
|
||||
}
|
||||
for i in range(device_count)
|
||||
]
|
||||
except NVMLError as error:
|
||||
print("Error fetching GPU information using nvml: %s", error)
|
||||
return None
|
||||
|
||||
monitor = MemoryMonitor(False)
|
||||
memory_before_test = monitor.measure_gpu_usage()
|
||||
|
||||
if start_memory is None:
|
||||
start_memory = memory_before_test
|
||||
if start_memory is None:
|
||||
return None
|
||||
if func is None:
|
||||
return start_memory
|
||||
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
with ThreadPoolExecutor() as executor:
|
||||
monitor = MemoryMonitor()
|
||||
mem_thread = executor.submit(monitor.measure_gpu_usage)
|
||||
try:
|
||||
fn_thread = executor.submit(func)
|
||||
_ = fn_thread.result()
|
||||
finally:
|
||||
monitor.keep_measuring = False
|
||||
max_usage = mem_thread.result()
|
||||
|
||||
if max_usage is None:
|
||||
return None
|
||||
|
||||
print(f"GPU memory usage: before={memory_before_test} peak={max_usage}")
|
||||
if len(start_memory) >= 1 and len(max_usage) >= 1 and len(start_memory) == len(max_usage):
|
||||
# When there are multiple GPUs, we will check the one with maximum usage.
|
||||
max_used = 0
|
||||
for i, memory_before in enumerate(start_memory):
|
||||
before = memory_before["max_used_MB"]
|
||||
after = max_usage[i]["max_used_MB"]
|
||||
used = after - before
|
||||
max_used = max(max_used, used)
|
||||
return max_used
|
||||
return None
|
||||
|
||||
|
||||
def get_ort_pipeline(model_name: str, directory: str, provider: str, disable_safety_checker: bool):
|
||||
from diffusers import DPMSolverMultistepScheduler, OnnxStableDiffusionPipeline
|
||||
|
||||
|
|
@ -61,14 +152,25 @@ def get_ort_pipeline(model_name: str, directory: str, provider: str, disable_saf
|
|||
return pipe
|
||||
|
||||
|
||||
def get_torch_pipeline(model_name: str, disable_safety_checker: bool):
|
||||
def get_torch_pipeline(model_name: str, disable_safety_checker: bool, enable_torch_compile: bool, use_xformers: bool):
|
||||
from diffusers import DPMSolverMultistepScheduler, StableDiffusionPipeline
|
||||
from torch import channels_last, float16
|
||||
|
||||
pipe = StableDiffusionPipeline.from_pretrained(
|
||||
model_name, torch_dtype=float16, revision="fp16", use_auth_token=True
|
||||
).to("cuda")
|
||||
pipe = StableDiffusionPipeline.from_pretrained(model_name, torch_dtype=float16).to("cuda")
|
||||
|
||||
pipe.unet.to(memory_format=channels_last) # in-place operation
|
||||
|
||||
if use_xformers:
|
||||
pipe.enable_xformers_memory_efficient_attention()
|
||||
|
||||
if enable_torch_compile:
|
||||
import torch
|
||||
|
||||
pipe.unet = torch.compile(pipe.unet)
|
||||
pipe.vae = torch.compile(pipe.vae)
|
||||
pipe.text_encoder = torch.compile(pipe.text_encoder)
|
||||
print("Torch compiled unet, vae and text_encoder")
|
||||
|
||||
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
|
||||
pipe.set_progress_bar_config(disable=True)
|
||||
|
||||
|
|
@ -84,14 +186,21 @@ def get_image_filename_prefix(engine: str, model_name: str, batch_size: int, dis
|
|||
return f"{engine}_{short_model_name}_b{batch_size}" + ("" if disable_safety_checker else "_safe")
|
||||
|
||||
|
||||
def run_ort_pipeline(pipe, batch_size: int, image_filename_prefix: str, height, width, steps, num_prompts, batch_count):
|
||||
def run_ort_pipeline(
|
||||
pipe, batch_size: int, image_filename_prefix: str, height, width, steps, num_prompts, batch_count, start_memory
|
||||
):
|
||||
from diffusers import OnnxStableDiffusionPipeline
|
||||
|
||||
assert isinstance(pipe, OnnxStableDiffusionPipeline)
|
||||
|
||||
prompts = example_prompts()
|
||||
|
||||
pipe("warm up", height, width, num_inference_steps=steps)
|
||||
def warmup():
|
||||
pipe("warm up", height, width, num_inference_steps=steps, num_images_per_prompt=batch_size)
|
||||
|
||||
# Run warm up, and measure GPU memory of two runs (The first run has cuDNN algo search so it might need more memory)
|
||||
first_run_memory = measure_gpu_memory(warmup, start_memory)
|
||||
second_run_memory = measure_gpu_memory(warmup, start_memory)
|
||||
|
||||
latency_list = []
|
||||
for i, prompt in enumerate(prompts):
|
||||
|
|
@ -111,21 +220,42 @@ def run_ort_pipeline(pipe, batch_size: int, image_filename_prefix: str, height,
|
|||
inference_end = time.time()
|
||||
latency = inference_end - inference_start
|
||||
latency_list.append(latency)
|
||||
print(f"Inference took {latency} seconds")
|
||||
print(f"Inference took {latency:.3f} seconds")
|
||||
for k, image in enumerate(images):
|
||||
image.save(f"{image_filename_prefix}_{i}_{j}_{k}.jpg")
|
||||
|
||||
print("Average latency in seconds:", sum(latency_list) / len(latency_list))
|
||||
from onnxruntime import __version__ as ort_version
|
||||
|
||||
return {
|
||||
"engine": "onnxruntime",
|
||||
"version": ort_version,
|
||||
"height": height,
|
||||
"width": width,
|
||||
"steps": steps,
|
||||
"batch_size": batch_size,
|
||||
"batch_count": batch_count,
|
||||
"num_prompts": num_prompts,
|
||||
"average_latency": sum(latency_list) / len(latency_list),
|
||||
"median_latency": statistics.median(latency_list),
|
||||
"first_run_memory_MB": first_run_memory,
|
||||
"second_run_memory_MB": second_run_memory,
|
||||
}
|
||||
|
||||
|
||||
def run_torch_pipeline(
|
||||
pipe, batch_size: int, image_filename_prefix: str, height, width, steps, num_prompts, batch_count
|
||||
pipe, batch_size: int, image_filename_prefix: str, height, width, steps, num_prompts, batch_count, start_memory
|
||||
):
|
||||
import torch
|
||||
|
||||
prompts = example_prompts()
|
||||
|
||||
pipe("warm up", height, width, num_inference_steps=steps)
|
||||
# total 2 runs of warm up, and measure GPU memory
|
||||
def warmup():
|
||||
pipe("warm up", height, width, num_inference_steps=steps, num_images_per_prompt=batch_size)
|
||||
|
||||
# Run warm up, and measure GPU memory of two runs (The first run has cuDNN algo search so it might need more memory)
|
||||
first_run_memory = measure_gpu_memory(warmup, start_memory)
|
||||
second_run_memory = measure_gpu_memory(warmup, start_memory)
|
||||
|
||||
torch.set_grad_enabled(False)
|
||||
|
||||
|
|
@ -151,11 +281,24 @@ def run_torch_pipeline(
|
|||
inference_end = time.time()
|
||||
latency = inference_end - inference_start
|
||||
latency_list.append(latency)
|
||||
print(f"Inference took {latency} seconds")
|
||||
print(f"Inference took {latency:.3f} seconds")
|
||||
for k, image in enumerate(images):
|
||||
image.save(f"{image_filename_prefix}_{i}_{j}_{k}.jpg")
|
||||
|
||||
print("Average latency in seconds:", sum(latency_list) / len(latency_list))
|
||||
return {
|
||||
"engine": "torch",
|
||||
"version": torch.__version__,
|
||||
"height": height,
|
||||
"width": width,
|
||||
"steps": steps,
|
||||
"batch_size": batch_size,
|
||||
"batch_count": batch_count,
|
||||
"num_prompts": num_prompts,
|
||||
"average_latency": sum(latency_list) / len(latency_list),
|
||||
"median_latency": statistics.median(latency_list),
|
||||
"first_run_memory_MB": first_run_memory,
|
||||
"second_run_memory_MB": second_run_memory,
|
||||
}
|
||||
|
||||
|
||||
def run_ort(
|
||||
|
|
@ -169,6 +312,7 @@ def run_ort(
|
|||
steps,
|
||||
num_prompts,
|
||||
batch_count,
|
||||
start_memory,
|
||||
):
|
||||
load_start = time.time()
|
||||
pipe = get_ort_pipeline(model_name, directory, provider, disable_safety_checker)
|
||||
|
|
@ -176,18 +320,33 @@ def run_ort(
|
|||
print(f"Model loading took {load_end - load_start} seconds")
|
||||
|
||||
image_filename_prefix = get_image_filename_prefix("ort", model_name, batch_size, disable_safety_checker)
|
||||
run_ort_pipeline(pipe, batch_size, image_filename_prefix, height, width, steps, num_prompts, batch_count)
|
||||
result = run_ort_pipeline(
|
||||
pipe, batch_size, image_filename_prefix, height, width, steps, num_prompts, batch_count, start_memory
|
||||
)
|
||||
|
||||
result.update(
|
||||
{
|
||||
"model_name": model_name,
|
||||
"directory": directory,
|
||||
"provider": provider,
|
||||
"disable_safety_checker": disable_safety_checker,
|
||||
}
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
def run_torch(
|
||||
model_name: str,
|
||||
batch_size: int,
|
||||
disable_safety_checker: bool,
|
||||
enable_torch_compile: bool,
|
||||
use_xformers: bool,
|
||||
height,
|
||||
width,
|
||||
steps,
|
||||
num_prompts,
|
||||
batch_count,
|
||||
start_memory,
|
||||
):
|
||||
import torch
|
||||
|
||||
|
|
@ -198,13 +357,31 @@ def run_torch(
|
|||
torch.set_grad_enabled(False)
|
||||
|
||||
load_start = time.time()
|
||||
pipe = get_torch_pipeline(model_name, disable_safety_checker)
|
||||
pipe = get_torch_pipeline(model_name, disable_safety_checker, enable_torch_compile, use_xformers)
|
||||
load_end = time.time()
|
||||
print(f"Model loading took {load_end - load_start} seconds")
|
||||
|
||||
image_filename_prefix = get_image_filename_prefix("torch", model_name, batch_size, disable_safety_checker)
|
||||
with torch.inference_mode():
|
||||
run_torch_pipeline(pipe, batch_size, image_filename_prefix, height, width, steps, num_prompts, batch_count)
|
||||
|
||||
if not enable_torch_compile:
|
||||
with torch.inference_mode():
|
||||
result = run_torch_pipeline(
|
||||
pipe, batch_size, image_filename_prefix, height, width, steps, num_prompts, batch_count, start_memory
|
||||
)
|
||||
else:
|
||||
result = run_torch_pipeline(
|
||||
pipe, batch_size, image_filename_prefix, height, width, steps, num_prompts, batch_count, start_memory
|
||||
)
|
||||
|
||||
result.update(
|
||||
{
|
||||
"model_name": model_name,
|
||||
"directory": None,
|
||||
"provider": "compile" if enable_torch_compile else "xformers" if use_xformers else "default",
|
||||
"disable_safety_checker": disable_safety_checker,
|
||||
}
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
def parse_arguments():
|
||||
|
|
@ -246,6 +423,22 @@ def parse_arguments():
|
|||
)
|
||||
parser.set_defaults(enable_safety_checker=False)
|
||||
|
||||
parser.add_argument(
|
||||
"--enable_torch_compile",
|
||||
required=False,
|
||||
action="store_true",
|
||||
help="Enable compile unet for PyTorch 2.0",
|
||||
)
|
||||
parser.set_defaults(enable_torch_compile=False)
|
||||
|
||||
parser.add_argument(
|
||||
"--use_xformers",
|
||||
required=False,
|
||||
action="store_true",
|
||||
help="Use xformers for PyTorch",
|
||||
)
|
||||
parser.set_defaults(use_xformers=False)
|
||||
|
||||
parser.add_argument(
|
||||
"-b",
|
||||
"--batch_size",
|
||||
|
|
@ -307,19 +500,15 @@ def main():
|
|||
args = parse_arguments()
|
||||
print(args)
|
||||
|
||||
start_memory = measure_gpu_memory(None)
|
||||
print("GPU memory used before loading models:", start_memory)
|
||||
|
||||
sd_model = SD_MODELS[args.version]
|
||||
if args.engine == "onnxruntime":
|
||||
assert args.pipeline, "--pipeline should be specified for onnxruntime engine"
|
||||
|
||||
if args.batch_size > 1:
|
||||
# Need remove a line https://github.com/huggingface/diffusers/blob/a66f2baeb782e091dde4e1e6394e46f169e5ba58/src/diffusers/pipelines/stable_diffusion/pipeline_onnx_stable_diffusion.py#L307
|
||||
# in diffuers to run batch_size > 1.
|
||||
assert (
|
||||
not args.enable_safety_checker
|
||||
), "batch_size > 1 is not compatible with safety checker due to a bug in diffuers"
|
||||
|
||||
provider = "CUDAExecutionProvider" # TODO: use ["CUDAExecutionProvider", "CPUExecutionProvider"] in diffuers
|
||||
run_ort(
|
||||
result = run_ort(
|
||||
sd_model,
|
||||
args.pipeline,
|
||||
provider,
|
||||
|
|
@ -330,19 +519,52 @@ def main():
|
|||
args.steps,
|
||||
args.num_prompts,
|
||||
args.batch_count,
|
||||
start_memory,
|
||||
)
|
||||
else:
|
||||
run_torch(
|
||||
result = run_torch(
|
||||
sd_model,
|
||||
args.batch_size,
|
||||
not args.enable_safety_checker,
|
||||
args.enable_torch_compile,
|
||||
args.use_xformers,
|
||||
args.height,
|
||||
args.width,
|
||||
args.steps,
|
||||
args.num_prompts,
|
||||
args.batch_count,
|
||||
start_memory,
|
||||
)
|
||||
|
||||
print(result)
|
||||
|
||||
with open("benchmark_result.csv", mode="a", newline="") as csv_file:
|
||||
column_names = [
|
||||
"model_name",
|
||||
"directory",
|
||||
"engine",
|
||||
"version",
|
||||
"provider",
|
||||
"disable_safety_checker",
|
||||
"height",
|
||||
"width",
|
||||
"steps",
|
||||
"batch_size",
|
||||
"batch_count",
|
||||
"num_prompts",
|
||||
"average_latency",
|
||||
"median_latency",
|
||||
"first_run_memory_MB",
|
||||
"second_run_memory_MB",
|
||||
]
|
||||
csv_writer = csv.DictWriter(csv_file, fieldnames=column_names)
|
||||
csv_writer.writeheader()
|
||||
csv_writer.writerow(result)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
try:
|
||||
main()
|
||||
except Exception as e:
|
||||
tb = sys.exc_info()
|
||||
print(e.with_traceback(tb[2]))
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
# Install the following package in python 3.10
|
||||
diffusers==0.12.1
|
||||
diffusers==0.13.0
|
||||
transformers==4.26.0
|
||||
numpy==1.24.1
|
||||
accelerate==0.15.0
|
||||
|
|
@ -10,6 +10,7 @@ packaging==23.0
|
|||
protobuf==3.20.3
|
||||
psutil==5.9.4
|
||||
sympy==1.11.1
|
||||
py3nvml==0.2.7
|
||||
#Tested with PyTorch 1.13.1+cu117 (see pytorch.org for more download options).
|
||||
#--extra-index-url https://download.pytorch.org/whl/cu117
|
||||
#torch==1.13.1+cu117
|
||||
|
|
|
|||
Loading…
Reference in a new issue