diff --git a/psyneulink/library/compositions/autodiffcomposition.py b/psyneulink/library/compositions/autodiffcomposition.py index b859883fff..893fb1b17e 100644 --- a/psyneulink/library/compositions/autodiffcomposition.py +++ b/psyneulink/library/compositions/autodiffcomposition.py @@ -102,7 +102,7 @@ AutodiffComposition does not (currently) support the *automatic* construction of separate bias parameters. Thus, when constructing the PyTorch version of an AutodiffComposition, the `bias -`_ parameter of PyTorch modules are set to False. +`_ parameter of any PyTorch modules are set to False. However, biases can be implemented using `Composition_Bias_Nodes`. @@ -531,7 +531,7 @@ class AutodiffComposition(Composition): but slows performance (see `AutodiffComposition_PyTorch_LearningScale` for information about settings). synch_node_variables_with_torch : OPTIMIZATION_STEP, TRIAL, MINIBATCH, EPOCH, RUN or None - determines when to copy the current input to Pytorch nodes (modules) to the PsyNeuLink `variable + determines when to copy the current input to Pytorch functions to the PsyNeuLink `variable ` attribute of the corresponding PsyNeuLink `nodes `, if this is not specified in the call to `learn `. COMMENT: @@ -545,11 +545,11 @@ class AutodiffComposition(Composition): but slows performance (see `AutodiffComposition_PyTorch_LearningScale` for information about settings). synch_node_values_with_torch : OPTIMIZATION_STEP, MINIBATCH, EPOCH or RUN - determines when to copy the current output of Pytorch nodes (modules) to the PsyNeuLink `value - ` attribute of the corresponding PsyNeuLink `nodes `, if this is not - specified in the call to `learn `. Copying more frequently keeps the PsyNeuLink - representation more closely copying more frequently keeps them synchronized with parameter updates in Pytorch, - but slows performance (see `AutodiffComposition_PyTorch_LearningScale` for information about settings). + determines when to copy the current output of Pytorch functions to the PsyNeuLink `value ` + attribute of the corresponding PsyNeuLink `nodes `, if this is not specified in the call to + `learn `. Copying more frequently keeps the PsyNeuLink representation more closely + copying more frequently keeps them synchronized with parameter updates in Pytorch, but slows performance + (see `AutodiffComposition_PyTorch_LearningScale` for information about settings). synch_results_with_torch : OPTIMIZATION_STEP, TRIAL, MINIBATCH, EPOCH or RUN determines when to copy the current outputs of Pytorch nodes to the PsyNeuLink `results diff --git a/psyneulink/library/compositions/pytorchwrappers.py b/psyneulink/library/compositions/pytorchwrappers.py index 9e95a04032..60a6d4ad6d 100644 --- a/psyneulink/library/compositions/pytorchwrappers.py +++ b/psyneulink/library/compositions/pytorchwrappers.py @@ -58,12 +58,25 @@ class PytorchCompositionWrapper(torch.nn.Module): # # MODIFIED 7/29/24 NEW: NEEDED FOR torch MPS SUPPORT # class PytorchCompositionWrapper(torch.jit.ScriptModule): # MODIFIED 7/29/24 END - """Wrapper for a Composition as a Pytorch Module - Class that wraps a `Composition ` as a PyTorch module. + """Wrapper for a Composition as a Pytorch Module. + + Wraps an `AutodiffComposition` as a `PyTorch module + `_, with each `Mechanism ` in the + AutodiffComposition wrapped as a `PytorchMechanismWrapper`, each `Projection ` wrapped as a + `PytorchProjectionWrapper`, and any nested Compositions wrapped as `PytorchCompositionWrapper`\\s. Each + PytorchMechanismWrapper implements a Pytorch version of the `function(s) ` of the wrapped + `Mechanism`, which are executed in the PyTorchCompositionWrapper's `forward ` + method in the order specified by the AutodiffComposition's `scheduler `. The + `matrix ` Parameters of each wrapped `Projection` are assigned as parameters of the + `PytorchMechanismWrapper` Pytorch module and used, together with a Pytorch `matmul + `_ operation, to generate the input to each + PyTorch function as specified by the `PytorchProjectionWrapper`\\'s `graph `. The graph + can be visualized using the AutodiffComposition's `show_graph ` method and setting its + *show_pytorch* argument to True (see `PytorchShowGraph` for additional information). Two main responsibilities: - 1) Set up parameters of PyTorch model & information required for forward computation: + 1) Set up functions and parameters of PyTorch module required for it forward computation: Handle nested compositions (flattened in infer_backpropagation_learning_pathways): Deal with Projections into and/or out of a nested Composition as shown in figure below: (note: Projections in outer Composition to/from a nested Composition's CIMs are learnable, @@ -115,12 +128,14 @@ class PytorchCompositionWrapper(torch.nn.Module): `AutodiffComposition` being wrapped. wrapped_nodes : List[PytorchMechanismWrapper] - list of nodes in the PytorchCompositionWrapper corresponding to PyTorch modules. Generally these are - `Mechanisms ` wrapped in a `PytorchMechanismWrapper`, however, if the `AutodiffComposition` - being wrapped is itself a nested Composition, then the wrapped nodes are `PytorchCompositionWrapper` objects. - When the PyTorch model is executed these are "flattened" into a single PyTorch module, which can be visualized - using the AutodiffComposition's `show_graph ` method and setting its *show_pytorch* - argument to True (see `PytorchShowGraph` for additional information). + list of nodes in the PytorchCompositionWrapper corresponding to the PyTorch functions that comprise the + forward method of the Pytorch module implemented by the PytorchCompositionWrapper. Generally these are + `Mechanisms ` wrapped in a `PytorchMechanismWrapper`, however, if the `AutodiffComposition` Node + being wrapped is a nested Composition, then the wrapped node is itself a `PytorchCompositionWrapper` object. + When the PyTorch model is executed, all of these are "flattened" into a single PyTorch module, corresponding + to the outermost AutodiffComposition being wrapped, which can be visualized using that AutodiffComposition's + `show_graph ` method and setting its *show_pytorch* argument to True (see + `PytorchShowGraph` for additional information). nodes_map : Dict[Node: PytorchMechanismWrapper or PytorchCompositionWrapper] maps psyneulink `Nodes ` to PytorchCompositionWrapper nodes. @@ -140,7 +155,7 @@ class PytorchCompositionWrapper(torch.nn.Module): assigned by AutodffComposition after the wrapper is created, which passes the parameters to the optimizer device : torch.device - device used to process torch Tensors in PyTorch modules + device used to process torch Tensors in PyTorch functions params : nn.ParameterList() list of PyTorch parameters (connection weight matrices) in the PyTorch model. @@ -857,7 +872,7 @@ def detach_all(self): class PytorchMechanismWrapper(): """Wrapper for a Mechanism in a PytorchCompositionWrapper - These comprise nodes of the PytorchCompositionWrapper, and generally correspond to modules of a Pytorch model. + These comprise nodes of the PytorchCompositionWrapper, and generally correspond to functions in a Pytorch model. Attributes ---------- @@ -1128,9 +1143,11 @@ def __repr__(self): class PytorchProjectionWrapper(): """Wrapper for Projection in a PytorchCompositionWrapper - The matrix of the wrapped `_projection ` corresponds to the parameters - (connection weights) of the PyTorch Module that is the `function ` of the - `receiver ` of the wrapped Projection. + The matrix of the wrapped `_projection ` is assigned as a parameter of + (set of connection weights in ) the PyTorch Module that, coupled with a corresponding input and `torch.matmul + `_ operation, provide the input to the Pytorch + function associated with the `Node ` of the AutdiffComposition that is the `receiver + ` of the wrapped Projection. .. note:: In the case of a nested Composition, the sender and/or receiver attributes may be mapped to different Node(s)