def run()

in optimum/commands/onnxruntime/optimize.py [0:0]

• 21 lines of code
• 8 McCabe index (conditional complexity)

    def run(self):
        from ...onnxruntime.configuration import AutoOptimizationConfig, ORTConfig
        from ...onnxruntime.optimization import ORTOptimizer

        if self.args.output == self.args.onnx_model:
            raise ValueError("The output directory must be different than the directory hosting the ONNX model.")

        save_dir = self.args.output

        file_names = [model.name for model in self.args.onnx_model.glob("*.onnx")]
        optimizer = ORTOptimizer.from_pretrained(self.args.onnx_model, file_names)

        if self.args.config:
            optimization_config = ORTConfig.from_pretrained(self.args.config).optimization
        elif self.args.O1:
            optimization_config = AutoOptimizationConfig.O1()
        elif self.args.O2:
            optimization_config = AutoOptimizationConfig.O2()
        elif self.args.O3:
            optimization_config = AutoOptimizationConfig.O3()
        elif self.args.O4:
            optimization_config = AutoOptimizationConfig.O4()
        else:
            raise ValueError("Either -O1, -O2, -O3, -O4 or -c must be specified.")

        optimizer.optimize(save_dir=save_dir, optimization_config=optimization_config)