Refactor ONNX & Olive.

lshqqytiger · Feb 7, 2024 · 7da1d6b · 7da1d6b
1 parent d500e58
commit 7da1d6b
Show file tree

Hide file tree

Showing 79 changed files with 3,279 additions and 4,114 deletions.
diff --git a/...s/olive_optimize/config_text_encoder.json → configs/olive/sd/text_encoder.json b/...s/olive_optimize/config_text_encoder.json → configs/olive/sd/text_encoder.json
@@ -2,9 +2,9 @@
   "input_model": {
     "type": "PyTorchModel",
     "config": {
-      "model_path": "runwayml/stable-diffusion-v1-5",
+      "model_path": "",
       "model_loader": "text_encoder_load",
-      "model_script": "modules/sd_olive_scripts.py",
+      "model_script": "modules/olive_script.py",
       "io_config": {
         "input_names": ["input_ids"],
         "output_names": ["last_hidden_state", "pooler_output"],
@@ -16,9 +16,7 @@
   "systems": {
     "local_system": {
       "type": "LocalSystem",
-      "config": {
-        "accelerators": ["gpu"]
-      }
+      "config": { "accelerators": ["gpu"] }
     }
   },
   "evaluators": {
@@ -29,7 +27,7 @@
           "type": "latency",
           "sub_types": [{ "name": "avg" }],
           "user_config": {
-            "user_script": "modules/sd_olive_scripts.py",
+            "user_script": "modules/olive_script.py",
             "dataloader_func": "text_encoder_data_loader",
             "batch_size": 1
           }
@@ -38,13 +36,18 @@
     }
   },
   "passes": {
-    "convert": {
-      "type": "OnnxConversion",
+    "optimize_CPUExecutionProvider": {
+      "type": "OrtTransformersOptimization",
+      "disable_search": true,
       "config": {
-        "target_opset": 14
+        "model_type": "clip",
+        "opt_level": 0,
+        "float16": false,
+        "use_gpu": false,
+        "keep_io_types": false
       }
     },
-    "optimize": {
+    "optimize_DmlExecutionProvider": {
       "type": "OrtTransformersOptimization",
       "disable_search": true,
       "config": {
@@ -75,12 +78,44 @@
           "group_norm_channels_last": false
         },
         "force_fp32_ops": ["RandomNormalLike"],
-        "force_fp16_inputs": {
-          "GroupNorm": [0, 1, 2]
-        }
+        "force_fp16_inputs": { "GroupNorm": [0, 1, 2] }
+      }
+    },
+    "optimize_CUDAExecutionProvider": {
+      "type": "OrtTransformersOptimization",
+      "disable_search": true,
+      "config": {
+        "model_type": "clip",
+        "opt_level": 0,
+        "float16": true,
+        "use_gpu": true,
+        "keep_io_types": false
+      }
+    },
+    "optimize_ROCMExecutionProvider": {
+      "type": "OrtTransformersOptimization",
+      "disable_search": true,
+      "config": {
+        "model_type": "clip",
+        "opt_level": 0,
+        "float16": true,
+        "use_gpu": true,
+        "keep_io_types": false
+      }
+    },
+    "quantization": {
+      "type": "OnnxDynamicQuantization",
+      "disable_search": true,
+      "config": {
+        "save_as_external_data": false,
+        "all_tensors_to_one_file": true,
+        "per_channel": false,
+        "reduce_range": false,
+        "MatMulConstBOnly": true
       }
     }
   },
+  "pass_flows": [["optimize_AutoExecutionProvider"]],
   "engine": {
     "search_strategy": {
       "execution_order": "joint",

diff --git a/configs/olive_optimize/config_unet.json → configs/olive/sd/unet.json b/configs/olive_optimize/config_unet.json → configs/olive/sd/unet.json
@@ -2,9 +2,9 @@
   "input_model": {
     "type": "PyTorchModel",
     "config": {
-      "model_path": "runwayml/stable-diffusion-v1-5",
+      "model_path": "",
       "model_loader": "unet_load",
-      "model_script": "modules/sd_olive_scripts.py",
+      "model_script": "modules/olive_script.py",
       "io_config": {
         "input_names": [
           "sample",
@@ -33,9 +33,7 @@
   "systems": {
     "local_system": {
       "type": "LocalSystem",
-      "config": {
-        "accelerators": ["gpu"]
-      }
+      "config": { "accelerators": ["gpu"] }
     }
   },
   "evaluators": {
@@ -46,7 +44,7 @@
           "type": "latency",
           "sub_types": [{ "name": "avg" }],
           "user_config": {
-            "user_script": "modules/sd_olive_scripts.py",
+            "user_script": "modules/olive_script.py",
             "dataloader_func": "unet_data_loader",
             "batch_size": 2
           }
@@ -55,16 +53,18 @@
     }
   },
   "passes": {
-    "convert": {
-      "type": "OnnxConversion",
+    "optimize_CPUExecutionProvider": {
+      "type": "OrtTransformersOptimization",
+      "disable_search": true,
       "config": {
-        "target_opset": 14,
-        "save_as_external_data": true,
-        "all_tensors_to_one_file": true,
-        "external_data_name": "weights.pb"
+        "model_type": "clip",
+        "opt_level": 0,
+        "float16": false,
+        "use_gpu": false,
+        "keep_io_types": false
       }
     },
-    "optimize": {
+    "optimize_DmlExecutionProvider": {
       "type": "OrtTransformersOptimization",
       "disable_search": true,
       "config": {
@@ -95,12 +95,44 @@
           "group_norm_channels_last": false
         },
         "force_fp32_ops": ["RandomNormalLike"],
-        "force_fp16_inputs": {
-          "GroupNorm": [0, 1, 2]
-        }
+        "force_fp16_inputs": { "GroupNorm": [0, 1, 2] }
+      }
+    },
+    "optimize_CUDAExecutionProvider": {
+      "type": "OrtTransformersOptimization",
+      "disable_search": true,
+      "config": {
+        "model_type": "unet",
+        "opt_level": 0,
+        "float16": true,
+        "use_gpu": true,
+        "keep_io_types": false
+      }
+    },
+    "optimize_ROCMExecutionProvider": {
+      "type": "OrtTransformersOptimization",
+      "disable_search": true,
+      "config": {
+        "model_type": "unet",
+        "opt_level": 0,
+        "float16": true,
+        "use_gpu": true,
+        "keep_io_types": false
+      }
+    },
+    "quantization": {
+      "type": "OnnxDynamicQuantization",
+      "disable_search": true,
+      "config": {
+        "save_as_external_data": false,
+        "all_tensors_to_one_file": true,
+        "per_channel": false,
+        "reduce_range": false,
+        "MatMulConstBOnly": true
       }
     }
   },
+  "pass_flows": [["optimize_AutoExecutionProvider"]],
   "engine": {
     "search_strategy": {
       "execution_order": "joint",

diff --git a/...gs/olive_optimize/config_vae_decoder.json → configs/olive/sd/vae_decoder.json b/...gs/olive_optimize/config_vae_decoder.json → configs/olive/sd/vae_decoder.json
@@ -2,9 +2,9 @@
   "input_model": {
     "type": "PyTorchModel",
     "config": {
-      "model_path": "runwayml/stable-diffusion-v1-5",
+      "model_path": "",
       "model_loader": "vae_decoder_load",
-      "model_script": "modules/sd_olive_scripts.py",
+      "model_script": "modules/olive_script.py",
       "io_config": {
         "input_names": ["latent_sample", "return_dict"],
         "output_names": ["sample"],
@@ -23,9 +23,7 @@
   "systems": {
     "local_system": {
       "type": "LocalSystem",
-      "config": {
-        "accelerators": ["gpu"]
-      }
+      "config": { "accelerators": ["gpu"] }
     }
   },
   "evaluators": {
@@ -36,7 +34,7 @@
           "type": "latency",
           "sub_types": [{ "name": "avg" }],
           "user_config": {
-            "user_script": "modules/sd_olive_scripts.py",
+            "user_script": "modules/olive_script.py",
             "dataloader_func": "vae_decoder_data_loader",
             "batch_size": 1
           }
@@ -45,13 +43,18 @@
     }
   },
   "passes": {
-    "convert": {
-      "type": "OnnxConversion",
+    "optimize_CPUExecutionProvider": {
+      "type": "OrtTransformersOptimization",
+      "disable_search": true,
       "config": {
-        "target_opset": 14
+        "model_type": "clip",
+        "opt_level": 0,
+        "float16": false,
+        "use_gpu": false,
+        "keep_io_types": false
       }
     },
-    "optimize": {
+    "optimize_DmlExecutionProvider": {
       "type": "OrtTransformersOptimization",
       "disable_search": true,
       "config": {
@@ -82,12 +85,44 @@
           "group_norm_channels_last": false
         },
         "force_fp32_ops": ["RandomNormalLike"],
-        "force_fp16_inputs": {
-          "GroupNorm": [0, 1, 2]
-        }
+        "force_fp16_inputs": { "GroupNorm": [0, 1, 2] }
+      }
+    },
+    "optimize_CUDAExecutionProvider": {
+      "type": "OrtTransformersOptimization",
+      "disable_search": true,
+      "config": {
+        "model_type": "vae",
+        "opt_level": 0,
+        "float16": true,
+        "use_gpu": true,
+        "keep_io_types": false
+      }
+    },
+    "optimize_ROCMExecutionProvider": {
+      "type": "OrtTransformersOptimization",
+      "disable_search": true,
+      "config": {
+        "model_type": "vae",
+        "opt_level": 0,
+        "float16": true,
+        "use_gpu": true,
+        "keep_io_types": false
+      }
+    },
+    "quantization": {
+      "type": "OnnxDynamicQuantization",
+      "disable_search": true,
+      "config": {
+        "save_as_external_data": false,
+        "all_tensors_to_one_file": true,
+        "per_channel": false,
+        "reduce_range": false,
+        "MatMulConstBOnly": true
       }
     }
   },
+  "pass_flows": [["optimize_AutoExecutionProvider"]],
   "engine": {
     "search_strategy": {
       "execution_order": "joint",