resnet_qat_lightning_module_cpu.json

{
    "input_model":{
        "type": "PyTorchModel",
        "config": {
            "model_path": "models/resnet_trained_for_cifar10.pt",
            "io_config": {
                "input_names": ["input"],
                "input_shapes": [[1, 3, 32, 32]],
                "output_names": ["output"],
                "dynamic_axes": {
                    "input": {"0": "batch_size"},
                    "output": {"0": "batch_size"}
                }
            }
        }
    },
    "evaluators": {
        "common_evaluator": {
            "metrics":[
                {
                    "name": "accuracy",
                    "type": "accuracy",
                    "sub_types": [{"name": "accuracy_score", "priority": 1}],
                    "user_config":{
                        "post_processing_func": "post_process",
                        "user_script": "user_script.py",
                        "dataloader_func": "create_dataloader",
                        "batch_size": 1
                    }
                },
                {
                    "name": "latency",
                    "type": "latency",
                    "sub_types": [{"name": "avg", "priority": 2}],
                    "user_config":{
                        "user_script": "user_script.py",
                        "dataloader_func": "create_dataloader",
                        "batch_size": 1
                    }
                }
            ]
        }
    },
    "passes": {
        "quantization_aware_training":{
            "type": "QuantizationAwareTraining",
            "config":{
                "user_script": "user_script.py",
                "num_epochs": 10,
                "ptl_data_module": "PTLDataModule",
                "ptl_module": "PTLModule",
                "pl_extra_args": {"accelerator": "gpu", "enable_checkpointing": true, "devices": 1},
                "modules_to_fuse": [["conv1", "bn1"], ["conv2", "bn2"], ["conv3", "bn3"]],
                "qconfig_func": "create_qat_config"
            }
        },
        "conversion": {
            "type": "OnnxConversion",
            "config": {
                "target_opset": 17
            }
        },
        "perf_tuning": {
            "type": "OrtPerfTuning",
            "config": {
                "user_script": "user_script.py",
                "dataloader_func": "create_dataloader",
                "batch_size": 1
            }
        }
    },
    "engine": {
        "search_strategy": {
            "execution_order": "pass-by-pass",
            "search_algorithm": "exhaustive"
        },
        "evaluator": "common_evaluator",
        "cache_dir": "cache",
        "output_dir": "models/resnet_qat_lightning_module_cpu"
    }
}