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 DataModel 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 
 21
 22
 23
 24
 25
 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
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 DataModel 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 = self.__class__.__doc__.strip().split("\n")[0].strip()
            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:
            schema = standardize_schema(data_model.schema())
            self._schema = schema
        else:
            self._schema = standardize_schema(schema)

        self.built = False
        self._variables = []
        self._tracker = Tracker(
            {
                "variables": (
                    lambda x: isinstance(x, backend.Variable),
                    self._variables,
                ),
            }
        )
        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

    def schema(self):
        return self._schema

    @property
    def variables(self):
        return self._variables[:]

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

    def _track_variable(self, variable):
        self._tracker.add_to_store("variables", variable)

    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):
        """Apply the backprop/optimization for each trainable variables
        that match the optimizer schema.
        """
        iteration = self._iteration.json().get("iteration")
        self._iteration.json().update({"iteration": iteration + 1})
        for variable in trainable_variables:
            if contains_schema(variable.schema(), self.schema()):
                await self.optimize(variable, reward=reward)

    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.schema(), self.schema()):
                await self.finalize(variable)

    async def optimize(self, trainable_variable, reward=None):
        """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) async

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

Source code in synalinks/src/optimizers/optimizer.py 
132
133
134
135
136
137
138
139
140
async def apply_optimization(self, trainable_variables, reward=None):
    """Apply the backprop/optimization for each trainable variables
    that match the optimizer schema.
    """
    iteration = self._iteration.json().get("iteration")
    self._iteration.json().update({"iteration": iteration + 1})
    for variable in trainable_variables:
        if contains_schema(variable.schema(), self.schema()):
            await self.optimize(variable, reward=reward)

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 
157
158
159
160
161
162
163
164
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 
142
143
144
145
146
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.schema(), self.schema()):
            await self.finalize(variable)

load_own_variables(store)

Set the state of this optimizer object.

Source code in synalinks/src/optimizers/optimizer.py 
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
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) 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 
148
149
150
151
152
153
154
155
async def optimize(self, trainable_variable, reward=None):
    """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 
101
102
103
104
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()