11972348325

Committed 22 Nov 2024 12:13PM UTC coverage: 83.758% (-0.06%) from 83.818%

Build # 11972348325

Build Type

Pull #3125

github

Committed by

jdcpni

Commit Message

Pull Request Pull Request #3125: patch/autodiff_pnl_showgraph

Run Details

9418 of 12474 branches covered (75.5%)

Branch coverage included in aggregate %.

5 of 9 new or added lines in 2 files covered. (55.56%)

264 existing lines in 4 files now uncovered.

32590 of 37680 relevant lines covered (86.49%)

0.86 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

20.65

/psyneulink/library/compositions/pytorchshowgraph.py

# Princeton University licenses this file to You under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.  You may obtain a copy of the License at:
#     http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software distributed under the License is distributed
# on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and limitations under the License.


# **************************************** PyTorch show_graph *********************************************************

from beartype import beartype

from psyneulink._typing import Optional, Union, Literal

from psyneulink.core.compositions import NodeRole
from psyneulink.core.compositions.showgraph import ShowGraph, SHOW_JUST_LEARNING_PROJECTIONS, SHOW_LEARNING
from psyneulink.core.components.mechanisms.mechanism import Mechanism
from psyneulink.core.components.mechanisms.processing.compositioninterfacemechanism import CompositionInterfaceMechanism
from psyneulink.core.components.mechanisms.modulatory.control.controlmechanism import ControlMechanism
from psyneulink.core.components.projections.projection import Projection
from psyneulink.core.components.projections.pathway.mappingprojection import MappingProjection
from psyneulink.core.components.projections.modulatory.controlprojection import ControlProjection
from psyneulink.core.llvm import ExecutionMode
from psyneulink.core.globals.context import ContextFlags, handle_external_context
from psyneulink.core.globals.keywords import BOLD, INSET, NESTED, PNL

__all__ = ['SHOW_PYTORCH']

SHOW_PYTORCH = 'show_pytorch'
EXCLUDE_FROM_GRADIENT_CALC_LINE_STYLE = 'exclude_from_gradient_calc_line_style'
EXCLUDE_FROM_GRADIENT_CALC_COLOR = 'exclude_from_gradient_calc_color'

class PytorchShowGraph(ShowGraph):
    """ShowGraph object with `show_graph <ShowGraph.show_graph>` method for displaying `Composition`.

    This is a subclass of the `ShowGraph` class that is used to display the graph of a `Composition` used for learning
    in `PyTorch mode <Composition_Learning_AutodiffComposition>` (also see `AutodiffComposition_PyTorch`).  In this mode,
    any `nested Compositions <AutodiffComposition_Nesting>` are "flattened" (i.e., incorporated into the outermost
    Composition); also, any `Nodes <Composition_Nodes>`` designated as `exclude_from_gradient_calc
    <PytorchMechanismWrapper.exclude_from_gradient_calc>` are moved to the end of the graph (as they are executed
    after the gradient calculation), and any Projections designated as `exclude_in_autodiff
    <Projection.exclude_in_autodiff>` are not shown as they are not used in the gradient calculations at all.

    Arguments
    ---------

    show_pytorch : keyword : default 'PYTORCH'
        specifies that the PyTorch version of the graph should be shown.

    """

    def __init__(self, *args, **kwargs):
        self.show_pytorch = kwargs.pop(SHOW_PYTORCH, False)
        super().__init__(*args, **kwargs)

    @beartype
    @handle_external_context(source=ContextFlags.COMPOSITION)
    def show_graph(self, *args, **kwargs):
        """Override of show_graph to check for autodiff-specific options
        If show_pytorch==True, build pytorch rep of autofiffcomposition
        If show_learning==PNL, infer backpropagation learning pathways for Python version of graph
        """
        if SHOW_LEARNING in kwargs and kwargs[SHOW_LEARNING] == PNL:
            self.composition.infer_backpropagation_learning_pathways(ExecutionMode.Python)
            kwargs[SHOW_LEARNING] = True
            return super().show_graph(*args, **kwargs)
        self.show_pytorch = kwargs.pop(SHOW_PYTORCH, self.show_pytorch)
        context = kwargs.get('context')
        if self.show_pytorch:
            self.pytorch_rep = self.composition._build_pytorch_representation(context)
        self.exclude_from_gradient_calc_line_style = kwargs.pop(EXCLUDE_FROM_GRADIENT_CALC_LINE_STYLE, 'dotted')
        self.exclude_from_gradient_calc_color = kwargs.pop(EXCLUDE_FROM_GRADIENT_CALC_COLOR, 'brown')
        return super().show_graph(*args, **kwargs)

    def _get_processing_graph(self, composition, context):
        """Helper method that creates dependencies graph for nodes of autodiffcomposition used in Pytorch mode"""
        if self.show_pytorch:
            processing_graph = {}
            projections = self._get_projections(composition, context)
            # 7/9/24 FIX: COULD DO THIS BY ITERATING OVER PROJECTIONS INSTEAD OF NODES
            for node in self._get_nodes(composition, context):
                dependencies = set()
                for projection in projections:
                    if node is projection.receiver.owner:
                        dependencies.add(projection.sender.owner)
                    # Add dependency of INPUT node of nested graph on node in outer graph that projects to it
                    elif (isinstance(projection.receiver.owner, CompositionInterfaceMechanism) and
                          projection.receiver.owner._get_destination_info_from_input_CIM(projection.receiver)[1]
                          is node):
                        dependencies.add(projection.sender.owner)
                processing_graph[node] = dependencies
            # Add TARGET nodes
            for node in self.composition.learning_components:
                processing_graph[node] = set([afferent.sender.owner for afferent in node.path_afferents])
            return processing_graph
        else:
            return super()._get_processing_graph(composition, context)

    def _get_nodes(self, composition, context):
        """Override to return nodes of PytorchCompositionWrapper rather than autodiffcomposition"""
        if self.show_pytorch:
            nodes = list(self.pytorch_rep.nodes_map.keys())
            return nodes
        else:
            return super()._get_nodes(composition, context)

    def _get_projections(self, composition, context):
        """Override to return nodes of Pytorch graph"""
        if self.show_pytorch:
            projections = list(self.pytorch_rep.projections_map.keys())
            # FIX: NEED TO ADD PROJECTIONS TO NESTED COMPS THAT ARE TO CIM
            # Add any Projections to TARGET nodes
            projections += [afferent
                            for node in self.composition.learning_components
                            for afferent in node.path_afferents
                            if not isinstance(afferent.sender.owner, CompositionInterfaceMechanism)]
            return projections
        else:
            return super()._get_projections(composition, context)

    def _proj_in_composition(self, proj, composition_projections, context)->bool:
        """Override to include direct Projections from outer to nested comps in Pytorch mode"""
        if self.show_pytorch:
            processing_graph = self._get_processing_graph(self.composition, context)
            if proj in composition_projections:
                return True
            # If proj is betw. a sender and receiver specified in the processing_graphl, then it is in the autodiffcomp
            elif (proj.receiver.owner in processing_graph
                  and proj.sender.owner in processing_graph[proj.receiver.owner]):
                return True
            else:
                return False
        else:
            return super()._proj_in_composition(proj, composition_projections, context)

    def _get_roles_by_node(self, composition, node, context):
        """Override in Pytorch mode to return NodeRole.INTERNAL for all nodes in nested compositions"""
        if self.show_pytorch:
            try:
                return composition.get_roles_by_node(node)
            except:
                return [NodeRole.INTERNAL]
        if self.show_pytorch and node not in self.composition.nodes:
            return [NodeRole.INTERNAL]
        else:
            return super()._get_roles_by_node(composition, node, context)

    def _get_nodes_by_role(self, composition, role, context):
        """Override in Pytorch mode to return all nodes in nested compositions as INTERNAL"""
        if self.show_pytorch and composition is not self.composition:
            return None
        else:
            return super()._get_nodes_by_role(composition, role, context)

    def _implement_graph_node(self, g, rcvr, context, *args, **kwargs):
        """Override to assign EXCLUDE_FROM_GRADIENT_CALC nodes their own style in Pytorch mode"""
        if self.show_pytorch:
            if self.pytorch_rep.nodes_map[rcvr].exclude_from_gradient_calc:
                kwargs['style'] = self.exclude_from_gradient_calc_line_style
                kwargs['color'] = self.exclude_from_gradient_calc_color
            g.node(*args, **kwargs)
        else:
            return super()._implement_graph_node( g, rcvr, context, *args, **kwargs)

    def _implement_graph_edge(self, graph, proj, context, *args, **kwargs):
        """Override to assign custom attributes to edges"""

        if self.show_pytorch:
            kwargs['color'] = self.default_node_color

            modulatory_node = None
            if proj.parameter_ports[0].mod_afferents:
                modulatory_node = self.pytorch_rep.nodes_map[proj.parameter_ports[0].mod_afferents[0].sender.owner]

            if proj in self.pytorch_rep.projections_map:

                # If Projection is a LearningProjection that is active, assign color and arrowhead of a LearningProjection
                if proj.learnable or self.pytorch_rep.projections_map[proj].matrix.requires_grad:
                    kwargs['color'] = self.learning_color

                # If Projection is from a ModulatoryMechanism that is excluded from gradient calculations, assign that style
                elif modulatory_node and modulatory_node.exclude_from_gradient_calc:
                    kwargs['color'] = self.exclude_from_gradient_calc_color
                    kwargs['style'] = self.exclude_from_gradient_calc_line_style

            elif self._proj_in_composition(proj, self.pytorch_rep.projections_map, context) and proj.learnable:
                kwargs['color'] = self.learning_color

            graph.edge(*args, **kwargs)

        else:
            return super()._implement_graph_edge(graph, proj, context, *args, **kwargs)

1	# Princeton University licenses this file to You under the Apache License, Version 2.0 (the "License");
2	# you may not use this file except in compliance with the License. You may obtain a copy of the License at:
3	# http://www.apache.org/licenses/LICENSE-2.0
4	# Unless required by applicable law or agreed to in writing, software distributed under the License is distributed
5	# on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
6	# See the License for the specific language governing permissions and limitations under the License.
7
8
9	# ************************************** PyTorch show_graph *******************************************************
10
11	from beartype import beartype	1✔
12
13	from psyneulink._typing import Optional, Union, Literal	1✔
14
15	from psyneulink.core.compositions import NodeRole	1✔
16	from psyneulink.core.compositions.showgraph import ShowGraph, SHOW_JUST_LEARNING_PROJECTIONS, SHOW_LEARNING	1✔
17	from psyneulink.core.components.mechanisms.mechanism import Mechanism	1✔
18	from psyneulink.core.components.mechanisms.processing.compositioninterfacemechanism import CompositionInterfaceMechanism	1✔
19	from psyneulink.core.components.mechanisms.modulatory.control.controlmechanism import ControlMechanism	1✔
20	from psyneulink.core.components.projections.projection import Projection	1✔
21	from psyneulink.core.components.projections.pathway.mappingprojection import MappingProjection	1✔
22	from psyneulink.core.components.projections.modulatory.controlprojection import ControlProjection	1✔
23	from psyneulink.core.llvm import ExecutionMode	1✔
24	from psyneulink.core.globals.context import ContextFlags, handle_external_context	1✔
25	from psyneulink.core.globals.keywords import BOLD, INSET, NESTED, PNL	1✔
26
27	__all__ = ['SHOW_PYTORCH']	1✔
28
29	SHOW_PYTORCH = 'show_pytorch'	1✔
30	EXCLUDE_FROM_GRADIENT_CALC_LINE_STYLE = 'exclude_from_gradient_calc_line_style'	1✔
31	EXCLUDE_FROM_GRADIENT_CALC_COLOR = 'exclude_from_gradient_calc_color'	1✔
32
33	class PytorchShowGraph(ShowGraph):	1✔
34	"""ShowGraph object with `show_graph <ShowGraph.show_graph>` method for displaying `Composition`.
35
36	This is a subclass of the `ShowGraph` class that is used to display the graph of a `Composition` used for learning
37	in `PyTorch mode <Composition_Learning_AutodiffComposition>` (also see `AutodiffComposition_PyTorch`). In this mode,
38	any `nested Compositions <AutodiffComposition_Nesting>` are "flattened" (i.e., incorporated into the outermost
39	Composition); also, any `Nodes <Composition_Nodes>`` designated as `exclude_from_gradient_calc
40	<PytorchMechanismWrapper.exclude_from_gradient_calc>` are moved to the end of the graph (as they are executed
41	after the gradient calculation), and any Projections designated as `exclude_in_autodiff
42	<Projection.exclude_in_autodiff>` are not shown as they are not used in the gradient calculations at all.
43
44	Arguments
45	---------
46
47	show_pytorch : keyword : default 'PYTORCH'
48	specifies that the PyTorch version of the graph should be shown.
49
50	"""
51
52	def __init__(self, args, *kwargs):	1✔
53	self.show_pytorch = kwargs.pop(SHOW_PYTORCH, False)	1✔
54	super().__init__(args, *kwargs)	1✔
55
56	@beartype	1✔
57	@handle_external_context(source=ContextFlags.COMPOSITION)	1✔
58	def show_graph(self, args, *kwargs):	1✔
59	"""Override of show_graph to check for autodiff-specific options
60	If show_pytorch==True, build pytorch rep of autofiffcomposition
61	If show_learning==PNL, infer backpropagation learning pathways for Python version of graph
62	"""
NEW 63	if SHOW_LEARNING in kwargs and kwargs[SHOW_LEARNING] == PNL:	×
NEW 64	self.composition.infer_backpropagation_learning_pathways(ExecutionMode.Python)	×
NEW 65	kwargs[SHOW_LEARNING] = True	×
NEW 66	return super().show_graph(args, *kwargs)	×
67	self.show_pytorch = kwargs.pop(SHOW_PYTORCH, self.show_pytorch)	×
68	context = kwargs.get('context')	×
69	if self.show_pytorch:	×
70	self.pytorch_rep = self.composition._build_pytorch_representation(context)	×
71	self.exclude_from_gradient_calc_line_style = kwargs.pop(EXCLUDE_FROM_GRADIENT_CALC_LINE_STYLE, 'dotted')	×
72	self.exclude_from_gradient_calc_color = kwargs.pop(EXCLUDE_FROM_GRADIENT_CALC_COLOR, 'brown')	×
73	return super().show_graph(args, *kwargs)	×
74
75	def _get_processing_graph(self, composition, context):	1✔
76	"""Helper method that creates dependencies graph for nodes of autodiffcomposition used in Pytorch mode"""
77	if self.show_pytorch:	×
78	processing_graph = {}	×
79	projections = self._get_projections(composition, context)	×
80	# 7/9/24 FIX: COULD DO THIS BY ITERATING OVER PROJECTIONS INSTEAD OF NODES
81	for node in self._get_nodes(composition, context):	×
82	dependencies = set()	×
83	for projection in projections:	×
84	if node is projection.receiver.owner:	×
85	dependencies.add(projection.sender.owner)	×
86	# Add dependency of INPUT node of nested graph on node in outer graph that projects to it
87	elif (isinstance(projection.receiver.owner, CompositionInterfaceMechanism) and	×
88	projection.receiver.owner._get_destination_info_from_input_CIM(projection.receiver)[1]
89	is node):
90	dependencies.add(projection.sender.owner)	×
91	processing_graph[node] = dependencies	×
92	# Add TARGET nodes
93	for node in self.composition.learning_components:	×
94	processing_graph[node] = set([afferent.sender.owner for afferent in node.path_afferents])	×
95	return processing_graph	×
96	else:
97	return super()._get_processing_graph(composition, context)	×
98
99	def _get_nodes(self, composition, context):	1✔
100	"""Override to return nodes of PytorchCompositionWrapper rather than autodiffcomposition"""
101	if self.show_pytorch:	×
102	nodes = list(self.pytorch_rep.nodes_map.keys())	×
103	return nodes	×
104	else:
105	return super()._get_nodes(composition, context)	×
106
107	def _get_projections(self, composition, context):	1✔
108	"""Override to return nodes of Pytorch graph"""
109	if self.show_pytorch:	×
110	projections = list(self.pytorch_rep.projections_map.keys())	×
111	# FIX: NEED TO ADD PROJECTIONS TO NESTED COMPS THAT ARE TO CIM
112	# Add any Projections to TARGET nodes
113	projections += [afferent	×
114	for node in self.composition.learning_components
115	for afferent in node.path_afferents
116	if not isinstance(afferent.sender.owner, CompositionInterfaceMechanism)]
117	return projections	×
118	else:
119	return super()._get_projections(composition, context)	×
120
121	def _proj_in_composition(self, proj, composition_projections, context)->bool:	1✔
122	"""Override to include direct Projections from outer to nested comps in Pytorch mode"""
123	if self.show_pytorch:	×
124	processing_graph = self._get_processing_graph(self.composition, context)	×
125	if proj in composition_projections:	×
126	return True	×
127	# If proj is betw. a sender and receiver specified in the processing_graphl, then it is in the autodiffcomp
128	elif (proj.receiver.owner in processing_graph	×
129	and proj.sender.owner in processing_graph[proj.receiver.owner]):
130	return True	×
131	else:
132	return False	×
133	else:
134	return super()._proj_in_composition(proj, composition_projections, context)	×
135
136	def _get_roles_by_node(self, composition, node, context):	1✔
137	"""Override in Pytorch mode to return NodeRole.INTERNAL for all nodes in nested compositions"""
138	if self.show_pytorch:	×
139	try:	×
140	return composition.get_roles_by_node(node)	×
141	except:	×
142	return [NodeRole.INTERNAL]	×
143	if self.show_pytorch and node not in self.composition.nodes:	×
144	return [NodeRole.INTERNAL]	×
145	else:
146	return super()._get_roles_by_node(composition, node, context)	×
147
148	def _get_nodes_by_role(self, composition, role, context):	1✔
149	"""Override in Pytorch mode to return all nodes in nested compositions as INTERNAL"""
150	if self.show_pytorch and composition is not self.composition:	×
151	return None	×
152	else:
153	return super()._get_nodes_by_role(composition, role, context)	×
154
155	def _implement_graph_node(self, g, rcvr, context, args, *kwargs):	1✔
156	"""Override to assign EXCLUDE_FROM_GRADIENT_CALC nodes their own style in Pytorch mode"""
157	if self.show_pytorch:	×
158	if self.pytorch_rep.nodes_map[rcvr].exclude_from_gradient_calc:	×
159	kwargs['style'] = self.exclude_from_gradient_calc_line_style	×
160	kwargs['color'] = self.exclude_from_gradient_calc_color	×
161	g.node(args, *kwargs)	×
162	else:
163	return super()._implement_graph_node( g, rcvr, context, args, *kwargs)	×
164
165	def _implement_graph_edge(self, graph, proj, context, args, *kwargs):	1✔
166	"""Override to assign custom attributes to edges"""
167
168	if self.show_pytorch:	×
169	kwargs['color'] = self.default_node_color	×
170
171	modulatory_node = None	×
172	if proj.parameter_ports[0].mod_afferents:	×
173	modulatory_node = self.pytorch_rep.nodes_map[proj.parameter_ports[0].mod_afferents[0].sender.owner]	×
174
175	if proj in self.pytorch_rep.projections_map:	×
176
177	# If Projection is a LearningProjection that is active, assign color and arrowhead of a LearningProjection
178	if proj.learnable or self.pytorch_rep.projections_map[proj].matrix.requires_grad:	×
179	kwargs['color'] = self.learning_color	×
180
181	# If Projection is from a ModulatoryMechanism that is excluded from gradient calculations, assign that style
182	elif modulatory_node and modulatory_node.exclude_from_gradient_calc:	×
183	kwargs['color'] = self.exclude_from_gradient_calc_color	×
184	kwargs['style'] = self.exclude_from_gradient_calc_line_style	×
185
186	elif self._proj_in_composition(proj, self.pytorch_rep.projections_map, context) and proj.learnable:	×
187	kwargs['color'] = self.learning_color	×
188
189	graph.edge(args, *kwargs)	×
190
191	else:
192	return super()._implement_graph_edge(graph, proj, context, args, *kwargs)	×

PrincetonUniversity / PsyNeuLink / 11972348325

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous