ZO training support

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "DeepQuant"]
+	path = DeepQuant
+	url = https://github.com/JanCSEM/DeepQuant.git

Onnx4Deeploy.py

@@ -34,13 +34,15 @@ def list_available_models():
         CCTExporter,
         EpiDeNetExporter,
         LightweightCnnExporter,
+        QLiteCnnExporter,
         MambaExporter,
         MIBMInetExporter,
         MobileNetV2Exporter,
         MobileViTExporter,
         ResNetExporter,
         SimpleMlpExporter,
         SleepConViTExporter,
+        QSleepConViTExporter
     )
 
     models = {
@@ -140,6 +142,18 @@ def list_available_models():
             "input_shape": "(B, 1, 28, 28)",
             "classes": 10,
         },
+        "QLiteCNN": {
+            "class": QLiteCnnExporter,
+            "description": "QLite CNN (Compact CNN for image classification)",
+            "input_shape": "(B, 1, 28, 28)",
+            "classes": 10,
+        },
+        "QSleepConViT": {
+            "class": QSleepConViTExporter,
+            "description": "QLite SleepConViT (Quantized Vision Transformer for Sleep Stage Classification)",
+            "input_shape": "(B, 1, 3000)",
+            "classes": 5
+        }
     }
     return models
 
@@ -155,6 +169,7 @@ def list_available_operators():
         # Matrix operations
         "Gemm": "General matrix multiplication",
         "MatMul": "Matrix multiplication",
+        "Conv2d": "2D convolution",
         # Pooling
         "MaxPool": "Max pooling",
         "AveragePool": "Average pooling",
@@ -169,10 +184,20 @@ def list_available_operators():
         "ConvGradX": "Convolution input gradient",
         "ConvGradW": "Convolution weight gradient",
         "ConvGradB": "Convolution bias gradient",
+        # ZO
+        "PerturbNormal": "Perturb input with gaussian random noise",
+        "PerturbUniform": "Perturb input with uniform random noise",
+        "PerturbTriangle": "Perturb input with triangle random noise",
+        "PerturbRademacher": "Perturb input with Rademacher random noise",
+        "PerturbEggroll": "Perturb input with Eggroll random noise",
+        "RQSPerturbrademacher": "Perturb input with quantized Rademacher random noise",
+        "RQSPerturbUniform": "Perturb input with quantized Uniform random noise",
+
         # Others
         "ReduceSum": "Sum reduction",
         "SoftmaxCrossEntropy": "Softmax cross entropy",
         "ReluGrad": "ReLU gradient",
+
     }
     return operators
 
@@ -248,7 +273,7 @@ def generate_operator(operator_name: str, output_path: Optional[str] = None):
         sys.exit(1)
 
 
-def generate_model(model_name: str, mode: str, output_path: Optional[str] = None):
+def generate_model(model_name: str, mode: str, output_path: Optional[str] = None, noise_type: str = "gaussian"):
     """Generate model ONNX"""
     print(f"\n{'='*70}")
     print(f"🚀 Generating model: {model_name} ({mode.upper()} mode)")
@@ -292,12 +317,21 @@ def generate_model(model_name: str, mode: str, output_path: Optional[str] = None
         if mode == "infer":
             onnx_file = exporter.export_inference()
             mode_desc = "Inference mode"
+        elif mode == "q-infer":
+            onnx_file = exporter.export_inference(quant=True)
+            mode_desc = "Quantized Inference mode"
         elif mode == "train":
             onnx_file = exporter.export_training()
             mode_desc = "Training mode"
+        elif mode == "zo-train":
+            onnx_file = exporter.export_zo_training(noise_type=noise_type)
+            mode_desc = "Zeroth-order Training mode"
+        elif mode == "q-zo-train":
+            onnx_file = exporter.export_zo_training(noise_type=noise_type, quant=True)
+            mode_desc = "Quantized Zeroth-order Training mode"
         else:
             print(f"❌ Unknown mode: {mode}")
-            print("   Available modes: infer, train")
+            print("   Available modes: infer, train, zo-train, q-infer, q-zo-train")
             sys.exit(1)
 
         print(f"\n{'='*70}")
@@ -311,7 +345,9 @@ def generate_model(model_name: str, mode: str, output_path: Optional[str] = None
         files_to_check = ["network.onnx", "inputs.npz", "outputs.npz"]
         if mode == "train":
             files_to_check.extend(["network_train.onnx", "optimizer_model.onnx"])
-
+        elif mode in ["zo-train", "q-zo-train"]:
+            files_to_check.append("network_zo.onnx")
+
         for file in files_to_check:
             file_path = output_dir / file
             if file_path.exists():
@@ -406,9 +442,9 @@ def main():
         "-mode",
         "--mode",
         type=str,
-        choices=["infer", "train"],
+        choices=["infer", "train", "zo-train", "q-infer", "q-zo-train"],
         default="infer",
-        help="Model export mode: infer (inference) or train (training) [default: infer]",
+        help="Model export mode: infer (inference), train (BP training), zo-train (zeroth-order training), q-infer (quantized inference), or q-zo-train (quantized zeroth-order training) [default: infer]",
     )
 
     # Output path
@@ -429,7 +465,8 @@ def main():
 
     # Other options
     parser.add_argument("--examples", action="store_true", help="Show usage examples")
-
+    parser.add_argument("--noise-type", type=str, choices=["gaussian", "uniform", "triangle", "rademacher", "eggroll", "rqs_rademacher", "rqs_uniform"], 
+                        default="gaussian", help="Noise type for perturbation operators [default: gaussian]")
     # Parse arguments
     args = parser.parse_args()
 
@@ -472,7 +509,7 @@ def main():
     if args.operator:
         generate_operator(args.operator, args.output)
     elif args.model:
-        generate_model(args.model, args.mode, args.output)
+        generate_model(args.model, args.mode, args.output, args.noise_type)
 
 
 if __name__ == "__main__":

gen_noise_tests.sh

@@ -0,0 +1,7 @@
+#!/bin/bash
+
+python3 Onnx4Deeploy.py -operator PerturbNormal -o PerturbNormal
+python3 Onnx4Deeploy.py -operator PerturbUniform -o PerturbUniform
+python3 Onnx4Deeploy.py -operator PerturbRademacher -o PerturbRademacher
+python3 Onnx4Deeploy.py -operator PerturbTriangle -o PerturbTriangle
+python3 Onnx4Deeploy.py -operator PerturbEggroll -o PerturbEggroll

gen_zo_model_tests.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+
+# LiteCNN
+python3 Onnx4Deeploy.py -model LightweightCNN -mode infer -o LiteCNN
+python3 Onnx4Deeploy.py -model LightweightCNN -mode zo-train -o LiteCNN-Rad --noise-type rademacher
+python3 Onnx4Deeploy.py -model LightweightCNN -mode zo-train -o LiteCNN-Lorp --noise-type eggroll
+python3 Onnx4Deeploy.py -model LightweightCNN -mode zo-train -o LiteCNN-Uniform --noise-type uniform
+python3 Onnx4Deeploy.py -model LightweightCNN -mode zo-train -o LiteCNN-Gaussian --noise-type gaussian
+
+# QLiteCNN
+python3 Onnx4Deeploy.py -model QLiteCNN -mode q-infer -o QLiteCNN
+python3 Onnx4Deeploy.py -model QLiteCNN -mode q-zo-train -o QLiteCNN-Rad --noise-type rademacher
+python3 Onnx4Deeploy.py -model QLiteCNN -mode q-zo-train -o QLiteCNN-Lorp --noise-type eggroll
+python3 Onnx4Deeploy.py -model QLiteCNN -mode q-zo-train -o QLiteCNN-Uniform --noise-type uniform
+python3 Onnx4Deeploy.py -model QLiteCNN -mode q-zo-train -o QLiteCNN-Gaussian --noise-type gaussian