saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Doc for Stable Diffusion CUDA Optimizations (#14830)

Browse source 
Add document for stable diffusion optimizations and benchmark.
This commit is contained in:
Tianlei Wu 2023-03-01 19:29:30 -08:00 committed by GitHub
parent c6074f3a4b
commit c66af46fc1
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 277 additions and 40 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

19
onnxruntime/python/tools/transformers/fusion_options.py
View file

@ -52,6 +52,7 @@ class FusionOptions:
# options for stable diffusion
if model_type in ["unet", "vae", "clip"]:
self.enable_nhwc_conv = True
self.enable_group_norm = True
self.enable_bias_splitgelu = True
self.enable_packed_qkv = True
@ -100,12 +101,18 @@ class FusionOptions:
options.disable_attention_mask()
if args.model_type in ["unet", "vae", "clip"]:
if args.disable_nhwc_conv:
options.enable_nhwc_conv = False
if args.disable_group_norm:
options.enable_group_norm = False
if args.disable_packed_kv:
options.enable_packed_kv = False
if args.disable_bias_splitgelu:
options.enable_bias_splitgelu = False
if args.disable_packed_qkv:
options.enable_packed_qkv = False
if args.disable_packed_kv:
options.enable_packed_kv = False
if args.disable_bias_add:
options.enable_bias_add = False
return options
@ -264,3 +271,11 @@ class FusionOptions:
help="not fuse BiasSplitGelu. Only works for model_type=unet",
)
parser.set_defaults(disable_bias_splitgelu=False)
parser.add_argument(
"--disable_nhwc_conv",
required=False,
action="store_true",
help="Do not use NhwcConv. Only works for model_type=unet or vae",
)
parser.set_defaults(disable_nhwc_conv=False)

225
onnxruntime/python/tools/transformers/models/stable_diffusion/README.md Normal file
View file

@ -0,0 +1,225 @@
# Stable Diffusion CUDA Optimization
## Overview
[Stable Diffusion](https://stability.ai/blog/stable-diffusion-announcement) is a text-to-image latent diffusion model for image generation. Explanation of the Stable Diffusion can be found in [Stable Diffusion with Diffusers](https://huggingface.co/blog/stable_diffusion).
## CUDA Optimizations for Stable Diffusion
ONNX Runtime uses the following optimizations to speed up Stable Diffusion in CUDA:
* [Flash Attention](https://arxiv.org/abs/2205.14135) for float16 precision. Flash Attention uses tiling to reduce number of GPU memory reads/writes, and improves performance with less memory for long sequence length. The kernel requires GPUs of Compute Capability >= 7.5 (like T4, A100, and RTX 2060~4090).
* [Memory Efficient Attention](https://arxiv.org/abs/2112.05682v2) for float32 precision or older GPUs (like V100). We used the fused multi-head attention kernel in CUTLASS, and the kernel was contributed by xFormers.
* Channel-last (NHWC) convolution. For NVidia GPU with Tensor Cores support, NHWC tensor layout is recommended for convolution. See [Tensor Layouts In Memory: NCHW vs NHWC](https://docs.nvidia.com/deeplearning/performance/dl-performance-convolutional/index.html#tensor-layout).
* GroupNorm kernel for NHWC tensor layout.
* SkipLayerNormalization which fuses LayerNormalization with Add bias and residual inputs.
* BiasSplitGelu is a fusion of Add bias with SplitGelu activation.
* BiasAdd fuses Add bias and residual.
* Reduce Transpose nodes by graph transformation.
Some CUDA kernels (Flash Attention, GroupNorm, SplitGelu and BiasAdd etc.) were originally implemented in TensorRT by Nvidia.
Compare to TensorRT, ONNX Runtime has some advantages on stable diffusion: (1) Supports older GPUs like V100. (2) Uses less GPU memory. (3) Supports float32 models and Stable Diffusion 2.* models.
To show the impact of each optimization on latency and GPU memory, we did an experiment on RTX 3060 GPU:
| Optimizations | Average Latency (batch_size=1) | Memory in MB (batch_size=1) | Average Latency (batch_size=8) | Memory in MB (batch_size=8) |
| ---------------------------------------------------------------------------------- | ------------------------------ | --------------------------- | ------------------------------ | --------------------------- |
| Raw FP32 models | 25.6 | 10,667 | OOM | OOM |
| FP16 baseline | 10.2 | 10,709 | OOM | OOM |
| FP16 baseline + FMHA | 6.1 | 7,719 | 39.1 | 10,821 |
| FP16 baseline + FMHA + NhwcConv | 5.5 | 7,656 | 38.8 | 11,615 |
| FP16 baseline + FMHA + NhwcConv + GroupNorm | 5.1 | 6,673 | 35.8 | 10,763 |
| FP16 baseline + FMHA + NhwcConv + GroupNorm + BiasSplitGelu | 4.9 | 4,447 | 33.7 | 6,669 |
| FP16 baseline + FMHA + NhwcConv + GroupNorm + BiasSplitGelu + Packed QKV | 4.8 | 4,625 | 33.5 | 6,663 |
| FP16 baseline + FMHA + NhwcConv + GroupNorm + BiasSplitGelu + Packed QKV + BiasAdd | 4.7 | 4,480 | 33.3 | 6,499 |
FP16 baseline contains optimizations available in ONNX Runtime 1.13 including LayerNormalization, SkipLayerNormalization, Gelu and float16 conversion.
Here FMHA means Attention and MultiHeadAttention operators with Flash Attention and Memory Efficient Attention kernels but inputs are not packed. Packed QKV means the inputs are packed.
The last two optimizations (Packed QKV and BiasAdd) are only available in nightly package. Compared to 1.14.1, nightly package has slight improvement in performance.
## Scripts:
| Script | Description |
| ---------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------- |
| [optimize_pipeline.py](https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/python/tools/transformers/models/stable_diffusion/optimize_pipeline.py) | Optimize Stable Diffusion ONNX models |
| [benchmark.py](https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/python/tools/transformers/models/stable_diffusion/benchmark.py) | Benchmark latency and memory of OnnxRuntime, xFormers or PyTorch 2.0 on stable diffusion. |
In below example, we run the scripts in source code directory. You can get source code like the following:
```
git clone https://github.com/microsoft/onnxruntime
cd onnxruntime/python/tools/transformers/models/stable_diffusion
```
## Example of Stable Diffusion 1.5
Below is an example to optimize Stable Diffusion 1.5 in Linux. For Windows OS, please change the format of path to be like `.\sd-v1-5` instead of `./sd-v1-5`.
### Setup Environment
First, Let's create a python environment using [AnaConda](https://www.anaconda.com/products/distribution#Downloads), then install packages in [requirements.txt](https://raw.githubusercontent.com/microsoft/onnxruntime/main/onnxruntime/python/tools/transformers/models/stable_diffusion/requirements.txt):
```
conda create -n py310 python=3.10
conda activate py310
pip install -r requirements.txt
```
For Windows, torch installed from pypi is CPU only. Need install PyTorch 1.13.1+cu117 instead like the following to support GPU:
```
pip install torch==1.13.1+cu117 --extra-index-url https://download.pytorch.org/whl/cu117
```
ONNX Runtime requires CUDA and [cuDNN](https://developer.nvidia.com/rdp/cudnn-download) for GPU inference. See https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html for compatible versions (like [CUDA 11.7](https://developer.nvidia.com/cuda-11-7-0-download-archive) and cuDNN 8.5.0.96 in Windows).
#### Install Nightly (Optional)
Skip this step if you use onnxruntime-gpu 1.14.* release package.
To try latest optimizations, you can install [ort-nightly-gpu](https://aiinfra.visualstudio.com/PublicPackages/_artifacts/feed/ORT-Nightly/PyPI/ort-nightly-gpu/) package like the following:
```
pip uninstall onnxruntime-gpu
pip install ort-nightly-gpu -i https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ORT-Nightly/pypi/simple/
```
The diffusers need a patch to work with ort-nightly-gpu: add `"ort-nightly-gpu",` below [this line](https://github.com/huggingface/diffusers/blob/1586186eea36e6a9b6f16c90ee32921c316c4b57/src/diffusers/utils/import_utils.py#L155).
### Export ONNX pipeline
This step will export stable diffusion 1.5 to ONNX model in float32 using script from diffusers. Before running the script, you need to be logged in via `huggingface-cli login`.
```
curl https://raw.githubusercontent.com/huggingface/diffusers/v0.13.0/scripts/convert_stable_diffusion_checkpoint_to_onnx.py > convert_sd_onnx.py
python convert_sd_onnx.py --model_path runwayml/stable-diffusion-v1-5 --output_path ./sd-v1-5
```
### Optimize ONNX Pipeline
Example to optimize the exported float32 ONNX models, and save to float16 models:
```
python optimize_pipeline.py -i ./sd-v1-5 -o ./sd-v1-5-fp16 --float16
```
If you installed ONNX Runtime v1.14, some optimizations (packed QKV and BiasAdd) will be disabled automatically since they are not available in v1.14.
For Stable Diffusion 2.1 model, you will need force Attention to run in float32 to avoid black image by appending `--force_fp32_ops unet:Attention` to the command line. If you are using nightly package, append `--force_fp32_ops unet:MultiHeadAttention` instead.
### Run Benchmark
The benchmark.py script will run a warm-up prompt twice, and measure the peak GPU memory usage in these two runs, then record them as first_run_memory_MB and second_run_memory_MB. Then it will run 5 runs to get average latency (in seconds), and output the results to benchmark_result.csv.
Note that the first run might need more time and memory: For example, cuDNN convolution algorithm search or model compile happens in the first run.
Example to benchmark the optimized pipeline with batch size 1:
```
python benchmark.py -p ./sd-v1-5-fp16/ -b 1
```
The default parameters are stable diffusion version=1.5, height=512, width=512, steps=50, batch_count=5. Run `python benchmark.py --help` for more information.
### Run Benchmark with xFormers
Run PyTorch 1.13.1+cu117 with xFormers like the following
```
python benchmark.py -e torch -b 1 --use_xformers
```
### Run Benchmark with PyTorch 2.0 with torch.compile
Let's create a new environment to run PyTorch 2.0:
```
conda create -n pt2 python=3.10
conda activate pt2
pip3 install numpy --pre torch --force-reinstall --extra-index-url https://download.pytorch.org/whl/nightly/cu117
pip install -r requirements.txt
python benchmark.py -e torch -b 1 --enable_torch_compile
```
Sometime, it complains ptxas not found when there are multiple CUDA versions installed. It can be fixed like `export TRITON_PTXAS_PATH=/usr/local/cuda-11.7/bin/ptxas` before running benchmark.
Note that torch.compile is not supported in Windows: we encountered error `Windows not yet supported for torch.compile`. So it is excluded from RTX 3060 results of Windows.
### Example Benchmark output
Common settings for below test results:
| model_name | disable_safety_checker | height | width | steps | batch_count | num_prompts |
| ------------------------------ | ---------------------- | ------ | ----- | ----- | ----------- | ----------- |
| runwayml/stable-diffusion-v1-5 | TRUE | 512 | 512 | 50 | 5 | 1 |
#### Results of RTX 3060 (in Windows 11)
| engine | version | provider | batch size | average latency | first run memory MB | second run memory MB |
| ----------- | ----------------------- | --------------------- | ---------- | --------------- | ------------------- | -------------------- |
| onnxruntime | 1.14.1 | CUDAExecutionProvider | 1 | 4.8 | 4,117 | 4,625 |
| torch | 2.0.0.dev20230220+cu117 | default | 1 | 5.6 | 4,330 | 4,050 |
| torch | 1.13.1+cu117 | xformers | 1 | 6.0 | 9,124 | 9,130 |
| onnxruntime | 1.14.1 | CUDAExecutionProvider | 4 | 17.7 | 6,659 | 6,659 |
| torch | 2.0.0.dev20230220+cu117 | default | 4 | 20.2 | 6,425 | 6,911 |
| torch | 1.13.1+cu117 | xformers | 4 | 21.6 | 10,407 | 10,409 |
| onnxruntime | 1.14.1 | CUDAExecutionProvider | 8 | 33.5 | 6,663 | 6,663 |
| torch | 2.0.0.dev20230220+cu117 | default | 8 | 39.8 | 10,894 | 10,782 |
| torch | 1.13.1+cu117 | xformers | 8 | 41.1 | 10,825 | 9,255 |
#### Results of V100-PCIE-16GB (in Ubuntu 20.04)
Results from Standard_NC6s_v3 Azure virtual machine:
| engine | version | provider | batch size | average latency | first run memory MB | second run memory MB |
| ----------- | ----------------------- | --------------------- | ---------- | --------------- | ------------------- | -------------------- |
| onnxruntime | 1.14.1 | CUDAExecutionProvider | 1 | 2.7 | 6,646 | 7,152 |
| torch | 2.0.0.dev20230220+cu117 | compile | 1 | 3.1 | 13,461 | 4,051 |
| torch | 2.0.0.dev20230220+cu117 | default | 1 | 2.7 | 13,461 | 4,041 |
| torch | 1.13.1+cu117 | xformers | 1 | 3.5 | 14,979 | 10,449 |
| onnxruntime | 1.14.1 | CUDAExecutionProvider | 4 | 8.4 | 7,114 | 7,114 |
| torch | 2.0.0.dev20230220+cu117 | compile | 4 | 8.0 | 14,015 | 7,085 |
| torch | 2.0.0.dev20230220+cu117 | default | 4 | 8.8 | 13,985 | 6,749 |
| torch | 1.13.1+cu117 | xformers | 4 | 9.1 | 12,969 | 8,421 |
| onnxruntime | 1.14.1 | CUDAExecutionProvider | 8 | 15.9 | 7,120 | 7,120 |
| torch | 2.0.0.dev20230220+cu117 | compile | 8 | 15.6 | 14,819 | 11,055 |
| torch | 2.0.0.dev20230220+cu117 | default | 8 | 16.9 | 14,603 | 10,563 |
| torch | 1.13.1+cu117 | xformers | 8 | 17.4 | 15,593 | 9,133 |
#### Results of T4 (in Ubuntu 20.04)
To make the result stable, we lock the frequency of T4 GPU like
`sudo nvidia-smi --lock-gpu-clocks=990` for fair comparison. See [nvidia blog](https://developer.nvidia.com/blog/advanced-api-performance-setstablepowerstate/) for more information. Note that performance might be slightly better without locking frequency.
Results are from Standard_NC4as_T4_v3 Azure virtual machine:
| engine | version | provider | batch size | average latency | first run memory MB | second run memory MB |
| ----------- | ----------------------- | --------------------- | ---------- | --------------- | ------------------- | -------------------- |
| onnxruntime | 1.14.1 | CUDAExecutionProvider | 1 | 5.6 | 4,925 | 4,925 |
| torch | 1.13.1+cu117 | xformers | 1 | 6.9 | 14,845 | 10,317 |
| torch | 2.0.0.dev20230226+cu117 | compile | 1 | 6.0 | 13,125 | 3,977 |
| torch | 2.0.0.dev20230226+cu117 | default | 1 | 6.3 | 13,127 | 3,979 |
| onnxruntime | 1.14.1 | CUDAExecutionProvider | 4 | 23.0 | 6,977 | 6,977 |
| torch | 1.13.1+cu117 | xformers | 4 | 25.8 | 12,819 | 8,269 |
| torch | 2.0.0.dev20230226+cu117 | compile | 4 | 22.1 | 14,751 | 6,697 |
| torch | 2.0.0.dev20230226+cu117 | default | 4 | 25.0 | 14,535 | 6,481 |
| onnxruntime | 1.14.1 | CUDAExecutionProvider | 8 | 46.4 | 6,779 | 6,779 |
| torch | 1.13.1+cu117 | xformers | 8 | 51.4 | 14,827 | 9,001 |
| torch | 2.0.0.dev20230226+cu117 | compile | 8 | 45.6 | 12,675 | 10,249 |
| torch | 2.0.0.dev20230226+cu117 | default | 8 | 50.5 | 12,077 | 9,653 |
### Future Works
There are other optimizations might improve the performance or reduce memory footprint:
* Use IO Binding in the pipeline. Currently the input and output of each model is in CPU, and extra data copy between GPU and CPU slows down the pipeline.
* Use CUDA graph to speed up inference.
* Export the whole pipeline into a single ONNX model. Currently, there are multiple ONNX models (CLIP, VAE and U-Net etc). Each model uses separated thread pool and memory allocator. Combine them into one model could share thread pool and memory allocator. The end result is more efficient and less memory footprint.
* For Stable Diffusion 2.1, we force Attention in fp32 to avoid black image. That slows down the inference significantly. We could potentially change attention kernel (like fp32 accumulation) to avoid the issue.
* Reduce GPU memory footprint by actively deleting buffers for intermediate results
* Reduce GPU memory footprint by CPU RAM Offloading
* Attention fusion in CLIP
* Safety Checker Optimization
* Leverage FP8 in latest GPU

24
onnxruntime/python/tools/transformers/models/stable_diffusion/benchmark.py
View file

@ -352,7 +352,6 @@ def run_torch(
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
# torch.backends.cuda.matmul.allow_tf32 = True
torch.set_grad_enabled(False)
@ -400,10 +399,11 @@ def parse_arguments():
parser.add_argument(
"-v",
"--version",
required=True,
required=False,
type=str,
choices=list(SD_MODELS.keys()),
help="Stable diffusion version like 1.5, 2.0 or 2.1",
default="1.5",
help="Stable diffusion version like 1.5, 2.0 or 2.1. Default is 1.5.",
)
parser.add_argument(
@ -444,8 +444,8 @@ def parse_arguments():
"--batch_size",
type=int,
default=1,
choices=[1, 2, 4, 8, 16, 32],
help="Number of images per batch",
choices=[1, 2, 4, 8, 10, 16, 32],
help="Number of images per batch. Default is 1.",
)
parser.add_argument(
@ -453,7 +453,7 @@ def parse_arguments():
required=False,
type=int,
default=512,
help="Output image height",
help="Output image height. Default is 512.",
)
parser.add_argument(
@ -461,7 +461,7 @@ def parse_arguments():
required=False,
type=int,
default=512,
help="Output image width",
help="Output image width. Default is 512.",
)
parser.add_argument(
@ -470,7 +470,7 @@ def parse_arguments():
required=False,
type=int,
default=50,
help="Number of steps",
help="Number of steps. Default is 50.",
)
parser.add_argument(
@ -479,7 +479,7 @@ def parse_arguments():
required=False,
type=int,
default=1,
help="Number of prompts",
help="Number of prompts. Default is 1.",
)
parser.add_argument(
@ -488,8 +488,8 @@ def parse_arguments():
required=False,
type=int,
choices=range(1, 11),
default=10,
help="Number of batches to test",
default=5,
help="Number of batches to test. Default is 5.",
)
args = parser.parse_args()
@ -507,7 +507,7 @@ def main():
if args.engine == "onnxruntime":
assert args.pipeline, "--pipeline should be specified for onnxruntime engine"
provider = "CUDAExecutionProvider" # TODO: use ["CUDAExecutionProvider", "CPUExecutionProvider"] in diffuers
provider = "CUDAExecutionProvider"
result = run_ort(
sd_model,
args.pipeline,

32
onnxruntime/python/tools/transformers/models/stable_diffusion/optimize_pipeline.py
View file

@ -5,25 +5,18 @@
#
# This script converts stable diffusion onnx models from float to half (mixed) precision for GPU inference.
#
# Before running this script, you need convert checkpoint to float32 onnx models like the following
# export ONNX_ROOT=./sd_onnx
# pip install -r requirements.txt
# huggingface-cli login
# wget https://raw.githubusercontent.com/huggingface/diffusers/v0.12.1/scripts/convert_stable_diffusion_checkpoint_to_onnx.py
# python convert_stable_diffusion_checkpoint_to_onnx.py --model_path runwayml/stable-diffusion-v1-5 --output_path $ONNX_ROOT/sd-v1-5
# python convert_stable_diffusion_checkpoint_to_onnx.py --model_path stabilityai/stable-diffusion-2-1 --output_path $ONNX_ROOT/sd-v2-1
# Note that this script might not be compatible with older or newer version of diffusers.
# Then you can use this script to convert them to float16 like the following:
# python optimize_pipeline.py -i $ONNX_ROOT/sd-v1-5 -o $ONNX_ROOT/sd-v1-5-fp16 --float16
# Or
# python -m onnxruntime.transformers.models.stable_diffusion.optimize_pipeline -i $ONNX_ROOT/sd-v1-5 -o $ONNX_ROOT/sd-v1-5-fp16 --float16
# Before running this script, follow README.md to setup python environment and convert stable diffusion checkpoint to float32 onnx models.
#
# Note that output model is for CUDA Execution Provider. It might not run in CPU Execution Provider.
# For example, the float32 ONNX pipeline is saved to ./sd-v1-5 directory, you can optimize and convert it to float16 like the following:
# python optimize_pipeline.py -i ./sd-v1-5 -o ./sd-v1-5-fp16 --float16
#
# Stable diffusion 2.1 model will get black images using float16 Attention. A walkaround is to force it in float32:
# python optimize_pipeline.py -i $ONNX_ROOT/sd-v2-1 -o $ONNX_ROOT/sd-v2-1-fp16 --float16 --force_fp32_ops unet:Attention
# Note that the optimized models are for CUDA Execution Provider. It might not run in other execution provider.
#
# Stable diffusion 2.1 model will get black images using float16 Attention. A walkaround is to force Attention to run in float32 like the following:
# python optimize_pipeline.py -i ./sd-v2-1 -o ./sd-v2-1-fp16 --float16 --force_fp32_ops unet:Attention
#
# If you are using nightly package (or built from source), you can force MultiHeadAttention to run in float32:
# python optimize_pipeline.py -i ./sd-v2-1 -o ./sd-v2-1-fp16 --float16 --force_fp32_ops unet:MultiHeadAttention
import argparse
import logging
@ -131,9 +124,10 @@ def optimize_sd_pipeline(
logger.info(f"Optimize {onnx_model_path}...")
fusion_options = FusionOptions(model_type)
if model_type in ["unet"]:
# There are some optimizations that are not available in v1.14 or older version
has_all_optimizations = version.parse(onnxruntime.__version__) > version.parse("1.14.0")
# Some optimizations are not available in v1.14 or older version: packed QKV and BiasAdd
has_all_optimizations = version.parse(onnxruntime.__version__) >= version.parse("1.15.0")
fusion_options.enable_packed_kv = float16
fusion_options.enable_packed_qkv = float16 and has_all_optimizations
fusion_options.enable_bias_add = has_all_optimizations

3
onnxruntime/python/tools/transformers/models/stable_diffusion/requirements.txt
View file

@ -11,6 +11,7 @@ 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).
xformers==0.0.16
#For Windows, need install PyTorch 1.13.1+cu117 since torch in pypi is CPU version
#--extra-index-url https://download.pytorch.org/whl/cu117
#torch==1.13.1+cu117

12
onnxruntime/python/tools/transformers/onnx_model_unet.py
View file

@ -89,7 +89,7 @@ class UnetOnnxModel(BertOnnxModel):
if total:
logger.info("Removed %d Transpose nodes", total)
def fuse_attention(self, options: Optional[FusionOptions] = None):
def fuse_multi_head_attention(self, options: Optional[FusionOptions] = None):
# Self Attention
enable_packed_qkv = (options is None) or options.enable_packed_qkv
self_attention_fusion = FusionAttentionUnet(
@ -139,7 +139,7 @@ class UnetOnnxModel(BertOnnxModel):
bias_split_gelu_fusion.apply()
if (options is None) or options.enable_attention:
self.fuse_attention()
self.fuse_multi_head_attention(options)
if (options is None) or options.enable_skip_layer_norm:
self.fuse_skip_layer_norm()
@ -149,8 +149,6 @@ class UnetOnnxModel(BertOnnxModel):
# Remove reshape nodes that having same shape of input and output based on symbolic shape inference.
self.utils.remove_useless_reshape_nodes()
self.convert_conv_to_nhwc()
if (options is None) or options.enable_bias_skip_layer_norm:
# Fuse SkipLayerNormalization and Add Bias before it.
self.fuse_add_bias_skip_layer_norm()
@ -158,7 +156,10 @@ class UnetOnnxModel(BertOnnxModel):
if options is not None and options.enable_gelu_approximation:
self.gelu_approximation()
self.merge_adjacent_transpose()
if options is None or options.enable_nhwc_conv:
self.convert_conv_to_nhwc()
self.merge_adjacent_transpose()
if options is not None and options.enable_bias_add:
self.fuse_bias_add()
@ -180,6 +181,7 @@ class UnetOnnxModel(BertOnnxModel):
"BiasSplitGelu",
"GroupNorm",
"NhwcConv",
"BiasAdd",
]
for op in ops:
nodes = self.get_nodes_by_op_type(op)

2
onnxruntime/python/tools/transformers/onnx_model_vae.py
View file

@ -19,7 +19,7 @@ class VaeOnnxModel(UnetOnnxModel):
assert (num_heads == 0 and hidden_size == 0) or (num_heads > 0 and hidden_size % num_heads == 0)
super().__init__(model, num_heads=num_heads, hidden_size=hidden_size)
def fuse_attention(self, options: Optional[FusionOptions] = None):
def fuse_multi_head_attention(self, options: Optional[FusionOptions] = None):
# Self Attention
self_attention_fusion = FusionAttentionVae(self, self.hidden_size, self.num_heads)
self_attention_fusion.apply()