|
| 1 | +from collections.abc import Callable |
| 2 | +import numpy as np |
| 3 | +from keras.saving import ( |
| 4 | + deserialize_keras_object as deserialize, |
| 5 | + register_keras_serializable as serializable, |
| 6 | + serialize_keras_object as serialize, |
| 7 | + get_registered_name, |
| 8 | + get_registered_object, |
| 9 | +) |
| 10 | +from .elementwise_transform import ElementwiseTransform |
| 11 | +from ...utils import filter_kwargs |
| 12 | +import inspect |
| 13 | + |
| 14 | + |
| 15 | +@serializable(package="bayesflow.adapters") |
| 16 | +class SerializableCustomTransform(ElementwiseTransform): |
| 17 | + """ |
| 18 | + Transforms a parameter using a pair of registered serializable forward and inverse functions. |
| 19 | +
|
| 20 | + Parameters |
| 21 | + ---------- |
| 22 | + serializable_forward_fn : function, no lambda |
| 23 | + Registered serializable function to transform the data in the forward pass. |
| 24 | + For the adapter to be serializable, this function has to be serializable |
| 25 | + as well (see Notes). Therefore, only proper functions and no lambda |
| 26 | + functions can be used here. |
| 27 | + serializable_inverse_fn : function, no lambda |
| 28 | + Function to transform the data in the inverse pass. |
| 29 | + For the adapter to be serializable, this function has to be serializable |
| 30 | + as well (see Notes). Therefore, only proper functions and no lambda |
| 31 | + functions can be used here. |
| 32 | +
|
| 33 | + Raises |
| 34 | + ------ |
| 35 | + ValueError |
| 36 | + When the provided functions are not registered serializable functions. |
| 37 | +
|
| 38 | + Notes |
| 39 | + ----- |
| 40 | + Important: The forward and inverse functions have to be registered with Keras. |
| 41 | + To do so, use the `@keras.saving.register_keras_serializable` decorator. |
| 42 | + They must also be registered (and identical) when loading the adapter |
| 43 | + at a later point in time. |
| 44 | +
|
| 45 | + """ |
| 46 | + |
| 47 | + def __init__( |
| 48 | + self, |
| 49 | + *, |
| 50 | + serializable_forward_fn: Callable[[np.ndarray, ...], np.ndarray], |
| 51 | + serializable_inverse_fn: Callable[[np.ndarray, ...], np.ndarray], |
| 52 | + ): |
| 53 | + super().__init__() |
| 54 | + |
| 55 | + self._check_serializable(serializable_forward_fn, label="serializable_forward_fn") |
| 56 | + self._check_serializable(serializable_inverse_fn, label="serializable_inverse_fn") |
| 57 | + self._forward = serializable_forward_fn |
| 58 | + self._inverse = serializable_inverse_fn |
| 59 | + |
| 60 | + @classmethod |
| 61 | + def _check_serializable(cls, function, label=""): |
| 62 | + GENERAL_EXAMPLE_CODE = f"""The example code below shows the structure of a correctly decorated function: |
| 63 | +
|
| 64 | +``` |
| 65 | +import keras |
| 66 | +
|
| 67 | +@keras.saving.register_keras_serializable('custom') |
| 68 | +def my_{label}(...): |
| 69 | + [your code goes here...] |
| 70 | +``` |
| 71 | +""" |
| 72 | + if function is None: |
| 73 | + raise TypeError( |
| 74 | + f"'{label}' must be a registered serializable function, was 'NoneType'.\n{GENERAL_EXAMPLE_CODE}" |
| 75 | + ) |
| 76 | + registered_name = get_registered_name(function) |
| 77 | + # check if function is a lambda function |
| 78 | + if registered_name == "<lambda>": |
| 79 | + raise ValueError( |
| 80 | + f"The provided function for '{label}' is a lambda function, " |
| 81 | + "which cannot be serialized. " |
| 82 | + "Please provide a registered serializable function by using the " |
| 83 | + "@keras.saving.register_keras_serializable decorator." |
| 84 | + f"\n{GENERAL_EXAMPLE_CODE}" |
| 85 | + ) |
| 86 | + if inspect.ismethod(function): |
| 87 | + raise ValueError( |
| 88 | + f"The provided value for '{label}' is a method, not a function. " |
| 89 | + "Methods cannot be serialized separately from their classes. " |
| 90 | + "Please provide a registered serializable function instead by " |
| 91 | + "moving the functionality to a function (i.e., outside of the class) and " |
| 92 | + "using the @keras.saving.register_keras_serializable decorator." |
| 93 | + f"\n{GENERAL_EXAMPLE_CODE}" |
| 94 | + ) |
| 95 | + registered_object_for_name = get_registered_object(registered_name) |
| 96 | + if registered_object_for_name is None: |
| 97 | + try: |
| 98 | + source_max_lines = 5 |
| 99 | + function_source_code = inspect.getsource(function).split("\n") |
| 100 | + if len(function_source_code) > source_max_lines: |
| 101 | + function_source_code = function_source_code[:source_max_lines] + [" [...]"] |
| 102 | + |
| 103 | + example_code = "For your provided function, this would look like this:\n\n" |
| 104 | + example_code += "\n".join( |
| 105 | + ["```", "import keras\n", "@keras.saving.register_keras_serializable('custom')"] |
| 106 | + + function_source_code |
| 107 | + + ["```"] |
| 108 | + ) |
| 109 | + except OSError: |
| 110 | + example_code = GENERAL_EXAMPLE_CODE |
| 111 | + raise ValueError( |
| 112 | + f"The provided function for '{label}' is not registered with Keras.\n" |
| 113 | + "Please register the function using the " |
| 114 | + "@keras.saving.register_keras_serializable decorator.\n" |
| 115 | + f"{example_code}" |
| 116 | + ) |
| 117 | + if registered_object_for_name is not function: |
| 118 | + raise ValueError( |
| 119 | + f"The provided function for '{label}' does not match the function " |
| 120 | + f"registered under its name '{registered_name}'. " |
| 121 | + f"(registered function: {registered_object_for_name}, provided function: {function}). " |
| 122 | + ) |
| 123 | + |
| 124 | + @classmethod |
| 125 | + def from_config(cls, config: dict, custom_objects=None) -> "SerializableCustomTransform": |
| 126 | + if get_registered_object(config["forward"]["config"], custom_objects) is None: |
| 127 | + provided_function_msg = "" |
| 128 | + if config["_forward_source_code"]: |
| 129 | + provided_function_msg = ( |
| 130 | + f"\nThe originally provided function was:\n\n```\n{config['_forward_source_code']}\n```" |
| 131 | + ) |
| 132 | + raise TypeError( |
| 133 | + "\n\nPLEASE READ HERE:\n" |
| 134 | + "-----------------\n" |
| 135 | + "The forward function that was provided as `serializable_forward_fn` " |
| 136 | + "is not registered with Keras, making deserialization impossible. " |
| 137 | + f"Please ensure that it is registered as '{config['forward']['config']}' and identical to the original " |
| 138 | + "function before loading your model." |
| 139 | + f"{provided_function_msg}" |
| 140 | + ) |
| 141 | + if get_registered_object(config["inverse"]["config"], custom_objects) is None: |
| 142 | + provided_function_msg = "" |
| 143 | + if config["_inverse_source_code"]: |
| 144 | + provided_function_msg = ( |
| 145 | + f"\nThe originally provided function was:\n\n```\n{config['_inverse_source_code']}\n```" |
| 146 | + ) |
| 147 | + raise TypeError( |
| 148 | + "\n\nPLEASE READ HERE:\n" |
| 149 | + "-----------------\n" |
| 150 | + "The inverse function that was provided as `serializable_inverse_fn` " |
| 151 | + "is not registered with Keras, making deserialization impossible. " |
| 152 | + f"Please ensure that it is registered as '{config['inverse']['config']}' and identical to the original " |
| 153 | + "function before loading your model." |
| 154 | + f"{provided_function_msg}" |
| 155 | + ) |
| 156 | + forward = deserialize(config["forward"], custom_objects) |
| 157 | + inverse = deserialize(config["inverse"], custom_objects) |
| 158 | + return cls( |
| 159 | + serializable_forward_fn=forward, |
| 160 | + serializable_inverse_fn=inverse, |
| 161 | + ) |
| 162 | + |
| 163 | + def get_config(self) -> dict: |
| 164 | + forward_source_code = inverse_source_code = None |
| 165 | + try: |
| 166 | + forward_source_code = inspect.getsource(self._forward) |
| 167 | + inverse_source_code = inspect.getsource(self._inverse) |
| 168 | + except OSError: |
| 169 | + pass |
| 170 | + return { |
| 171 | + "forward": serialize(self._forward), |
| 172 | + "inverse": serialize(self._inverse), |
| 173 | + "_forward_source_code": forward_source_code, |
| 174 | + "_inverse_source_code": inverse_source_code, |
| 175 | + } |
| 176 | + |
| 177 | + def forward(self, data: np.ndarray, **kwargs) -> np.ndarray: |
| 178 | + # filter kwargs so that other transform args like batch_size, strict, ... are not passed through |
| 179 | + kwargs = filter_kwargs(kwargs, self._forward) |
| 180 | + return self._forward(data, **kwargs) |
| 181 | + |
| 182 | + def inverse(self, data: np.ndarray, **kwargs) -> np.ndarray: |
| 183 | + kwargs = filter_kwargs(kwargs, self._inverse) |
| 184 | + return self._inverse(data, **kwargs) |
0 commit comments