Write trainables to module #470
Conversation
michaeldeistler
force-pushed
the
updatemodule
branch
from
ddbff57 to
80e1869
Compare
| @@ -1044,6 +1045,32 @@ def make_trainable( | |||
| f"Number of newly added trainable parameters: {num_created_parameters}. Total number of trainable parameters: {self.base.num_trainable_params}" | |||
| ) | |||
|
|
|||
| def write_trainables(self, trainable_params: List[Dict[str, jnp.ndarray]]): | |||
There was a problem hiding this comment.
Cool! This only works for Module tho, i.e. doing net.cell(0).write_trainables() will only write to self.base.nodes if self == self.base. To be able to call this from a View and for these changes to take affect for the module you need to write to self.base.nodes.loc[self._nodes_in_view], instead of self.nodes. See docstring of View. I think it might actually be nice if you could write only the trainables in View.
There was a problem hiding this comment.
doesn't work, see comments
There was a problem hiding this comment.
Feel free to merge this.
Regarding the special treatment of edges: since we are getting rid of synapse view down the line, we can treat edges like nodes. I think this would make a lot of things nicer. Currently working on this as part of #453.
* Write trainables to module * Allow to write the trainables into the module
This PR introduces a
write_trainables()which adds the trainables to the.nodesand.edgessuch that they can be used for plotting, printing,...Closes #356