Skip to content

Base Optimizer class

Bases: SynalinksSaveable

Optimizer base class: all Synalinks optimizers inherit from this class.

Parameters:

Name Type Description Default
schema dict

The schema of the variables that the optimizer can act upon.

 None
data_model DataModel

The backend data model that the optimizer can act upon, if no schema is specified, uses the data model to infer it.

 None
name str

The name of the optimizer.

 None
description str

The description of the optimizer.

 None
Source code in synalinks/src/optimizers/optimizer.py 
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
class Optimizer(SynalinksSaveable):
    """Optimizer base class: all Synalinks optimizers inherit from this class.

    Args:
        schema (dict): The schema of the variables that the optimizer can act upon.
        data_model (DataModel): The backend data model that the optimizer can act upon,
            if no schema is specified, uses the data model to infer it.
        name (str): The name of the optimizer.
        description (str): The description of the optimizer.
    """

    def __init__(
        self,
        schema=None,
        data_model=None,
        name=None,
        description=None,
        **kwargs,
    ):
        self._lock = False

        if kwargs:
            raise ValueError(f"Argument(s) not recognized: {kwargs}")

        if name is None:
            name = auto_name(self.__class__.__name__)
        self.name = name

        if description is None:
            if self.__class__.__doc__:
                description = docstring_parser.parse(
                    self.__class__.__doc__
                ).short_description
            else:
                description = ""
        self.description = description

        if not data_model and not schema:
            raise ValueError(
                "You should provide at least one argument "
                "between `data_model` or `schema`"
            )
        if not schema and data_model:
            schema = standardize_schema(data_model.get_schema())
        self._schema = schema

        self.built = False

        self._initialize_tracker()

        with backend.name_scope(self.name, caller=self):
            iterations = backend.Variable(
                initializer=Empty(data_model=Iteration),
                data_model=Iteration,
                trainable=False,
                name="iteration",
            )
        # self._track_variable(iterations)
        self._iteration = iterations

    @tracking.no_automatic_dependency_tracking
    def _initialize_tracker(self):
        if hasattr(self, "_tracker"):
            return

        trainable_variables = []
        non_trainable_variables = []
        modules = []
        self._tracker = tracking.Tracker(
            {
                "trainable_variables": (
                    lambda x: isinstance(x, backend.Variable) and x.trainable,
                    trainable_variables,
                ),
                "non_trainable_variables": (
                    lambda x: isinstance(x, backend.Variable) and not x.trainable,
                    non_trainable_variables,
                ),
                "modules": (
                    lambda x: isinstance(x, Module) and not isinstance(x, Metric),
                    modules,
                ),
            },
            exclusions={"non_trainable_variables": ["trainable_variables"]},
        )
        self._trainable_variables = trainable_variables
        self._non_trainable_variables = non_trainable_variables
        self._modules = modules

    def __setattr__(self, name, value):
        # Track Variables, Modules, Metrics.
        if name != "_tracker":
            if not hasattr(self, "_tracker"):
                self._initialize_tracker()
            value = self._tracker.track(value)
        return super().__setattr__(name, value)

    def get_schema(self):
        return self._schema

    @property
    def variables(self):
        return self._non_trainable_variables[:] + self._trainable_variables[:]

    @property
    def non_trainable_variables(self):
        return self._non_trainable_variables[:]

    @property
    def trainable_variables(self):
        variables = []
        for module in self._modules:
            variables.extend(module.trainable_variables)
        return variables

    @property
    def iterations(self):
        return self._iteration

    def save_own_variables(self, store):
        """Get the state of this optimizer object."""
        for i, variable in enumerate(self.variables):
            store[str(i)] = variable.numpy()

    def load_own_variables(self, store):
        """Set the state of this optimizer object."""
        if len(store.keys()) != len(self.variables):
            msg = (
                f"Skipping variable loading for optimizer '{self.name}', "
                f"because it has {len(self.variables)} variables whereas "
                f"the saved optimizer has {len(store.keys())} variables. "
            )
            if len(self.variables) == 0:
                msg += (
                    "This is likely because the optimizer has not been called/built yet."
                )
            warnings.warn(msg, stacklevel=2)
            return
        for i, variable in enumerate(self.variables):
            variable.assign(store[str(i)])

    def _check_super_called(self):
        if not hasattr(self, "_lock"):
            raise RuntimeError(
                f"In optimizer '{self.__class__.__name__}', you forgot to call "
                "`super().__init__()` as the first statement "
                "in the `__init__()` method. "
                "Go add it!"
            )

    async def apply_optimization(self, trainable_variables, reward=None, training=False):
        """Apply the backprop/optimization for each trainable variables
        that match the optimizer schema.
        """
        if not self.built:
            run_maybe_nested(self.build(trainable_variables))
        iteration = self._iteration.get("iteration")
        self._iteration.update({"iteration": iteration + 1})
        for variable in trainable_variables:
            if contains_schema(variable.get_schema(), self.get_schema()):
                await self.optimize(variable, reward=reward, training=training)

    async def finalize_variable_values(self, trainable_variables):
        """Finalize the optimization of the variables (cleanup/scaling etc.)."""
        for variable in trainable_variables:
            if contains_schema(variable.get_schema(), self.get_schema()):
                await self.finalize(variable)

    async def optimize(self, trainable_variable, reward=None, training=False):
        """Perform a backprop/optimization on a single variable.

        This function needs to be implemented by subclassed Optimizer.
        """
        raise NotImplementedError(
            "Optimizer subclasses must implement the `optimize()` method."
        )

    async def finalize(self, trainable_variable):
        """Finalize the optimization of the variable (cleanup/scaling etc.).

        This function needs to be implemented by subclassed Optimizer.
        """
        raise NotImplementedError(
            "Optimizer subclasses must implement the `finalize()` method."
        )

    def get_config(self):
        return {
            "name": self.name,
            "description": self.description,
            "schema": self.schema,
        }

    @classmethod
    def from_config(cls, config):
        return cls(**config)

    def __repr__(self):
        return f"<Optimizer name={self.name} description={self.description}>"

apply_optimization(trainable_variables, reward=None, training=False) async

Apply the backprop/optimization for each trainable variables that match the optimizer schema.

Source code in synalinks/src/optimizers/optimizer.py 
176
177
178
179
180
181
182
183
184
185
186
async def apply_optimization(self, trainable_variables, reward=None, training=False):
    """Apply the backprop/optimization for each trainable variables
    that match the optimizer schema.
    """
    if not self.built:
        run_maybe_nested(self.build(trainable_variables))
    iteration = self._iteration.get("iteration")
    self._iteration.update({"iteration": iteration + 1})
    for variable in trainable_variables:
        if contains_schema(variable.get_schema(), self.get_schema()):
            await self.optimize(variable, reward=reward, training=training)

finalize(trainable_variable) async

Finalize the optimization of the variable (cleanup/scaling etc.).

This function needs to be implemented by subclassed Optimizer.

Source code in synalinks/src/optimizers/optimizer.py 
203
204
205
206
207
208
209
210
async def finalize(self, trainable_variable):
    """Finalize the optimization of the variable (cleanup/scaling etc.).

    This function needs to be implemented by subclassed Optimizer.
    """
    raise NotImplementedError(
        "Optimizer subclasses must implement the `finalize()` method."
    )

finalize_variable_values(trainable_variables) async

Finalize the optimization of the variables (cleanup/scaling etc.).

Source code in synalinks/src/optimizers/optimizer.py 
188
189
190
191
192
async def finalize_variable_values(self, trainable_variables):
    """Finalize the optimization of the variables (cleanup/scaling etc.)."""
    for variable in trainable_variables:
        if contains_schema(variable.get_schema(), self.get_schema()):
            await self.finalize(variable)

load_own_variables(store)

Set the state of this optimizer object.

Source code in synalinks/src/optimizers/optimizer.py 
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
def load_own_variables(self, store):
    """Set the state of this optimizer object."""
    if len(store.keys()) != len(self.variables):
        msg = (
            f"Skipping variable loading for optimizer '{self.name}', "
            f"because it has {len(self.variables)} variables whereas "
            f"the saved optimizer has {len(store.keys())} variables. "
        )
        if len(self.variables) == 0:
            msg += (
                "This is likely because the optimizer has not been called/built yet."
            )
        warnings.warn(msg, stacklevel=2)
        return
    for i, variable in enumerate(self.variables):
        variable.assign(store[str(i)])

optimize(trainable_variable, reward=None, training=False) async

Perform a backprop/optimization on a single variable.

This function needs to be implemented by subclassed Optimizer.

Source code in synalinks/src/optimizers/optimizer.py 
194
195
196
197
198
199
200
201
async def optimize(self, trainable_variable, reward=None, training=False):
    """Perform a backprop/optimization on a single variable.

    This function needs to be implemented by subclassed Optimizer.
    """
    raise NotImplementedError(
        "Optimizer subclasses must implement the `optimize()` method."
    )

save_own_variables(store)

Get the state of this optimizer object.

Source code in synalinks/src/optimizers/optimizer.py 
145
146
147
148
def save_own_variables(self, store):
    """Get the state of this optimizer object."""
    for i, variable in enumerate(self.variables):
        store[str(i)] = variable.numpy()