-
Notifications
You must be signed in to change notification settings - Fork 234
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Adds ONNX Surrogate support from OMLT (#1308)
* fist commit * added tests * undid test issues * added direct load test * adde method tosave model to mirror keras Likely want to refactor this later * fixed acidental delet * Speling fixes and doc updates * clean up imports * Update pytest.ini * fix spelling in utility_minimi... * Fixed onnx dependancy * lint and rst fixes * fix gitignore * pylint and sphynix #2 * fixing rst * doc string fix * updates to comments * linting fixe * added onnx to distro --------- Co-authored-by: Kyle Skolfield <[email protected]> Co-authored-by: Keith Beattie <[email protected]> Co-authored-by: Ludovico Bianchi <[email protected]> Co-authored-by: Andrew Lee <[email protected]>
- Loading branch information
5 people
authored
Nov 4, 2024
1 parent
53f756b
commit 2e5a9a4
Showing
8 changed files
with
630 additions
and
108 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,187 @@ | ||
| ################################################################################# | ||
| # The Institute for the Design of Advanced Energy Systems Integrated Platform | ||
| # Framework (IDAES IP) was produced under the DOE Institute for the | ||
| # Design of Advanced Energy Systems (IDAES). | ||
| # | ||
| # Copyright (c) 2018-2023 by the software owners: The Regents of the | ||
| # University of California, through Lawrence Berkeley National Laboratory, | ||
| # National Technology & Engineering Solutions of Sandia, LLC, Carnegie Mellon | ||
| # University, West Virginia University Research Corporation, et al. | ||
| # All rights reserved. Please see the files COPYRIGHT.md and LICENSE.md | ||
| # for full copyright and license information. | ||
| ################################################################################# | ||
| """ | ||
| Interface for importing ONNX models into IDAES | ||
| """ | ||
| # TODO: Missing docstrings | ||
| # pylint: disable=missing-class-docstring | ||
| # pylint: disable=missing-function-docstring | ||
|
|
||
| from enum import Enum | ||
| import numpy as np | ||
|
|
||
| from pyomo.common.dependencies import attempt_import | ||
|
|
||
| from idaes.core.surrogate.base.surrogate_base import SurrogateBase | ||
| from idaes.core.surrogate.sampling.scaling import OffsetScaler | ||
|
|
||
| keras, keras_available = attempt_import("tensorflow.keras") | ||
| omlt, omlt_available = attempt_import("omlt") | ||
|
|
||
| if omlt_available: | ||
| from omlt import OmltBlock, OffsetScaling | ||
|
|
||
|
|
||
| class OMLTSurrogate(SurrogateBase): | ||
| def __init__( | ||
| self, | ||
| input_labels, | ||
| output_labels, | ||
| input_bounds, | ||
| input_scaler=None, | ||
| output_scaler=None, | ||
| ): | ||
| """ | ||
| Standard SurrogateObject for surrogates based on Keras models. | ||
| Utilizes the OMLT framework for importing Keras models to IDAES. | ||
| Contains methods to both populate a Pyomo Block with constraints | ||
| representing the surrogate and to evaluate the surrogate a set of user | ||
| provided points. | ||
| This constructor should only be used when first creating the surrogate within IDAES. | ||
| Once created, this object can be stored to disk using save_to_folder and loaded | ||
| with load_from_folder | ||
| Args: | ||
| onnx_model: Onnx model file to be loaded. | ||
| input_labels: list of str | ||
| The ordered list of labels corresponding to the inputs in the keras model | ||
| output_labels: list of str | ||
| The ordered list of labels corresponding to the outputs in the keras model | ||
| input_bounds: None of dict of tuples | ||
| Keys correspond to each of the input labels and values are the tuples of | ||
| bounds (lb, ub) | ||
| input_scaler: None or OffsetScaler | ||
| The scaler to be used for the inputs. If None, then no scaler is used | ||
| output_scaler: None of OffsetScaler | ||
| The scaler to be used for the outputs. If None, then no scaler is used | ||
| """ | ||
| super().__init__( | ||
| input_labels=input_labels, | ||
| output_labels=output_labels, | ||
| input_bounds=input_bounds, | ||
| ) | ||
|
|
||
| # make sure we are using the standard scaler | ||
| if ( | ||
| input_scaler is not None | ||
| and not isinstance(input_scaler, OffsetScaler) | ||
| or output_scaler is not None | ||
| and not isinstance(output_scaler, OffsetScaler) | ||
| ): | ||
| raise NotImplementedError("KerasSurrogate only supports the OffsetScaler.") | ||
|
|
||
| # check that the input labels match | ||
| if input_scaler is not None and input_scaler.expected_columns() != input_labels: | ||
| raise ValueError( | ||
| "KerasSurrogate created with input_labels that do not match" | ||
| " the expected columns in the input_scaler.\n" | ||
| "input_labels={}\n" | ||
| "input_scaler.expected_columns()={}".format( | ||
| input_labels, input_scaler.expected_columns() | ||
| ) | ||
| ) | ||
|
|
||
| # check that the output labels match | ||
| if ( | ||
| output_scaler is not None | ||
| and output_scaler.expected_columns() != output_labels | ||
| ): | ||
| raise ValueError( | ||
| "KerasSurrogate created with output_labels that do not match" | ||
| " the expected columns in the output_scaler.\n" | ||
| "output_labels={}\n" | ||
| "output_scaler.expected_columns()={}".format( | ||
| output_labels, output_scaler.expected_columns() | ||
| ) | ||
| ) | ||
|
|
||
| self._input_scaler = input_scaler | ||
| self._output_scaler = output_scaler | ||
|
|
||
| class Formulation(Enum): | ||
| FULL_SPACE = 1 | ||
| REDUCED_SPACE = 2 | ||
| RELU_BIGM = 3 | ||
| RELU_COMPLEMENTARITY = 4 | ||
|
|
||
| def generate_omlt_scaling_objecets(self): | ||
| offset_inputs = np.zeros(self.n_inputs()) | ||
| factor_inputs = np.ones(self.n_inputs()) | ||
| offset_outputs = np.zeros(self.n_outputs()) | ||
| factor_outputs = np.ones(self.n_outputs()) | ||
| if self._input_scaler: | ||
| offset_inputs = self._input_scaler.offset_series()[ | ||
| self.input_labels() | ||
| ].to_numpy() | ||
| factor_inputs = self._input_scaler.factor_series()[ | ||
| self.input_labels() | ||
| ].to_numpy() | ||
| if self._output_scaler: | ||
| offset_outputs = self._output_scaler.offset_series()[ | ||
| self.output_labels() | ||
| ].to_numpy() | ||
| factor_outputs = self._output_scaler.factor_series()[ | ||
| self.output_labels() | ||
| ].to_numpy() | ||
|
|
||
| omlt_scaling = OffsetScaling( | ||
| offset_inputs=offset_inputs, | ||
| factor_inputs=factor_inputs, | ||
| offset_outputs=offset_outputs, | ||
| factor_outputs=factor_outputs, | ||
| ) | ||
|
|
||
| # omlt takes *scaled* input bounds as a dictionary with int keys | ||
| input_bounds = dict(enumerate(self.input_bounds().values())) | ||
| scaled_input_bounds = omlt_scaling.get_scaled_input_expressions(input_bounds) | ||
| scaled_input_bounds = {i: tuple(bnd) for i, bnd in scaled_input_bounds.items()} | ||
| return omlt_scaling, scaled_input_bounds | ||
|
|
||
| def populate_block_with_net(self, block, formulation_object): | ||
| """ | ||
| Method to populate a Pyomo Block with the omlt model constraints and build its formulation. | ||
| Args: | ||
| block: Pyomo Block component | ||
| The block to be populated with variables and/or constraints. | ||
| formulation_object: omlt loaded network formulation | ||
| """ | ||
|
|
||
| block.nn = OmltBlock() | ||
| block.nn.build_formulation( | ||
| formulation_object, | ||
| ) | ||
|
|
||
| # input/output variables need to be constrained to be equal | ||
| # auto-created variables that come from OMLT. | ||
| input_idx_by_label = {s: i for i, s in enumerate(self._input_labels)} | ||
| input_vars_as_dict = block.input_vars_as_dict() | ||
|
|
||
| @block.Constraint(self._input_labels) | ||
| def input_surrogate_ties(m, input_label): | ||
| return ( | ||
| input_vars_as_dict[input_label] | ||
| == block.nn.inputs[input_idx_by_label[input_label]] | ||
| ) | ||
|
|
||
| output_idx_by_label = {s: i for i, s in enumerate(self._output_labels)} | ||
| output_vars_as_dict = block.output_vars_as_dict() | ||
|
|
||
| @block.Constraint(self._output_labels) | ||
| def output_surrogate_ties(m, output_label): | ||
| return ( | ||
| output_vars_as_dict[output_label] | ||
| == block.nn.outputs[output_idx_by_label[output_label]] | ||
| ) |
Oops, something went wrong.