7186854566

Committed 12 Dec 2023 08:20PM UTC coverage: 84.902% (+0.2%) from 84.676%

Build # 7186854566

Build Type

push

github

Committed by

web-flow

Commit Message

Merge branch 'devel' into models/ego/use_emcomposition

Run Details

13474 of 16603 branches covered (0.0%)

Branch coverage included in aggregate %.

263 of 266 new or added lines in 16 files covered. (98.87%)

188 existing lines in 9 files now uncovered.

31513 of 36384 relevant lines covered (86.61%)

0.87 hits per line

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

96.74

/psyneulink/library/compositions/pytorchEMcompositionwrapper.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.

# ********************************************* PytorchComponent *************************************************

"""PyTorch wrapper for EMComposition"""

import torch
# try:
#     import torch
# except (ImportError, ModuleNotFoundError):
#     torch = None
from typing import Optional

from psyneulink.library.compositions.pytorchwrappers import PytorchCompositionWrapper
from psyneulink.library.components.mechanisms.modulatory.learning.EMstoragemechanism import EMStorageMechanism

__all__ = ['PytorchEMCompositionWrapper']

class PytorchEMCompositionWrapper(PytorchCompositionWrapper):
    """Wrapper for EMComposition as a Pytorch Module"""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        # Assign storage_node (EMComposition's EMStorageMechanism)
        self.storage_node = [node for node in self.nodes_map.values()
                             if isinstance(node._mechanism, EMStorageMechanism)][0]

        # Get PytorchProjectionWrappers for Projections to match and retrieve nodes;
        #   used by get_memory() to construct memory_matrix and store_memory() to store entry in it
        pnl_storage_mech = self.storage_node._mechanism
        num_fields = len(pnl_storage_mech.input_ports)
        num_learning_signals = len(pnl_storage_mech.learning_signals)
        num_match_fields = num_learning_signals - num_fields

        # ProjectionWrappers for match nodes
        learning_signals_for_match_nodes = pnl_storage_mech.learning_signals[:num_match_fields]
        pnl_match_projs = [match_node_learning_signal.efferents[0].receiver.owner
                               for match_node_learning_signal in learning_signals_for_match_nodes]
        self.match_projection_wrappers = [self.projections_map[pnl_match_proj]
                           for pnl_match_proj in pnl_match_projs]

        # ProjectionWrappers for retrieve nodes
        learning_signals_for_retrieve_nodes = pnl_storage_mech.learning_signals[num_match_fields:]
        pnl_retrieve_projs = [retrieve_node_learning_signal.efferents[0].receiver.owner
                               for retrieve_node_learning_signal in learning_signals_for_retrieve_nodes]
        self.retrieve_projection_wrappers = [self.projections_map[pnl_retrieve_proj]
                           for pnl_retrieve_proj in pnl_retrieve_projs]

    def execute_node(self, node, variable, context):
        """Override to handle storage of entry to memory_matrix by EMStorage Function"""
        if node is self.storage_node:
            self.store_memory(variable, context)
        else:
            super().execute_node(node, variable, context)

    @property
    def memory(self)->Optional[torch.Tensor]:
        """Return list of memories in which rows (outer dimension) are memories for each field.
        These are derived from the matrix parameters of the afferent Projections to the retrieval_nodes
        """
        num_fields = len(self.storage_node.afferents)
        memory_matrices = [field.matrix for field in self.retrieve_projection_wrappers]
        memory_capacity = len(memory_matrices[0])
        return (None if not all(val for val in [num_fields, memory_matrices, memory_capacity])
                else torch.stack([torch.stack([memory_matrices[j][i]
                                               for j in range(num_fields)])
                                  for i in range(memory_capacity)]))

    def store_memory(self, memory_to_store, context):
        """Store variable in memory_matrix (parallel EMStorageMechanism._execute)

        For each node in query_input_nodes and value_input_nodes,
        assign its value to weights of corresponding afferents to corresponding match_node and/or retrieved_node.
        - memory = matrix of entries made up vectors for each field in each entry (row)
        - entry_to_store = query_input or value_input to store
        - field_projections = Projections the matrices of which comprise memory

        DIVISION OF LABOR between this method and function called by it
        store_memory (corresponds to EMStorageMechanism._execute)
         - compute norms to find weakest entry in memory
         - compute storage_prob to determine whether to store current entry in memory
         - call function with memory matrix for each field, to decay existing memory and assign input to weakest entry
        storage_node.function (corresponds to EMStorage._function):
         - decay existing memories
         - assign input to weakest entry (given index for passed from EMStorageMechanism)

        :return: List[2d tensor] updated memories
        """

        memory = self.memory
        assert memory is not None, f"PROGRAM ERROR: '{self.name}'.memory is None"

        # Get current parameter values from EMComposition's EMStorageMechanism
        mech = self.storage_node._mechanism
        random_state = mech.function.parameters.random_state._get(context)
        decay_rate = mech.parameters.decay_rate._get(context)      # modulable, so use getter
        storage_prob = mech.parameters.storage_prob._get(context)  # modulable, so use getter
        field_weights = mech.parameters.field_weights.get(context) # modulable, so use getter
        concatenation_node = mech.concatenation_node
        num_match_fields = 1 if concatenation_node else len([i for i in mech.field_types if i==1])

        # Find weakest memory (i.e., with lowest norm)
        field_norms = torch.empty((len(memory),len(memory[0])))
        for row in range(len(memory)):
            for col in range(len(memory[0])):
                field_norms[row][col] = torch.linalg.norm(memory[row][col])
        if field_weights is not None:
            field_norms *= field_weights
        row_norms = torch.sum(field_norms, axis=1)
        idx_of_weakest_memory = torch.argmin(row_norms)

        values = []
        for i, field_projection in enumerate(self.match_projection_wrappers + self.retrieve_projection_wrappers):
            if i < num_match_fields:
                # For match projections, get entry to store from value of sender of Projection matrix
                #   (this is to accomodate concatenation_node)
                axis = 0
                entry_to_store = field_projection.sender.value
                if concatenation_node is None:
                    assert (entry_to_store == memory_to_store[i]).all(), \
                        f"PROGRAM ERROR: misalignment between inputs and fields for storing them"
            else:
                # For retrieve projections, get entry to store from memory_to_store (which has inputs to all fields)
                axis = 1
                entry_to_store = memory_to_store[i - num_match_fields]
            # Get matrix containing memories for the field from the Projection
            field_memory_matrix = field_projection.matrix

            field_projection.matrix = self.storage_node.function(entry_to_store,
                                                                 memory_matrix=field_memory_matrix,
                                                                 axis=axis,
                                                                 storage_location=idx_of_weakest_memory,
                                                                 storage_prob=storage_prob,
                                                                 decay_rate=decay_rate,
                                                                 random_state=random_state)
            values.append(field_projection.matrix)

        self.storage_node.value = values
        return values

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	# ******************************************* PytorchComponent ***********************************************
9
10	"""PyTorch wrapper for EMComposition"""	1✔
11
12	import torch	1✔
13	# try:
14	# import torch
15	# except (ImportError, ModuleNotFoundError):
16	# torch = None
17	from typing import Optional	1✔
18
19	from psyneulink.library.compositions.pytorchwrappers import PytorchCompositionWrapper	1✔
20	from psyneulink.library.components.mechanisms.modulatory.learning.EMstoragemechanism import EMStorageMechanism	1✔
21
22	__all__ = ['PytorchEMCompositionWrapper']	1✔
23
24	class PytorchEMCompositionWrapper(PytorchCompositionWrapper):	1✔
25	"""Wrapper for EMComposition as a Pytorch Module"""
26
27	def __init__(self, args, *kwargs):	1✔
28	super().__init__(args, *kwargs)	1✔
29
30	# Assign storage_node (EMComposition's EMStorageMechanism)
31	self.storage_node = [node for node in self.nodes_map.values()	1✔
32	if isinstance(node._mechanism, EMStorageMechanism)][0]
33
34	# Get PytorchProjectionWrappers for Projections to match and retrieve nodes;
35	# used by get_memory() to construct memory_matrix and store_memory() to store entry in it
36	pnl_storage_mech = self.storage_node._mechanism	1✔
37	num_fields = len(pnl_storage_mech.input_ports)	1✔
38	num_learning_signals = len(pnl_storage_mech.learning_signals)	1✔
39	num_match_fields = num_learning_signals - num_fields	1✔
40
41	# ProjectionWrappers for match nodes
42	learning_signals_for_match_nodes = pnl_storage_mech.learning_signals[:num_match_fields]	1✔
43	pnl_match_projs = [match_node_learning_signal.efferents[0].receiver.owner	1✔
44	for match_node_learning_signal in learning_signals_for_match_nodes]
45	self.match_projection_wrappers = [self.projections_map[pnl_match_proj]	1✔
46	for pnl_match_proj in pnl_match_projs]
47
48	# ProjectionWrappers for retrieve nodes
49	learning_signals_for_retrieve_nodes = pnl_storage_mech.learning_signals[num_match_fields:]	1✔
50	pnl_retrieve_projs = [retrieve_node_learning_signal.efferents[0].receiver.owner	1✔
51	for retrieve_node_learning_signal in learning_signals_for_retrieve_nodes]
52	self.retrieve_projection_wrappers = [self.projections_map[pnl_retrieve_proj]	1✔
53	for pnl_retrieve_proj in pnl_retrieve_projs]
54
55	def execute_node(self, node, variable, context):	1✔
56	"""Override to handle storage of entry to memory_matrix by EMStorage Function"""
57	if node is self.storage_node:	1✔
58	self.store_memory(variable, context)	1✔
59	else:
60	super().execute_node(node, variable, context)	1✔
61
62	@property	1✔
63	def memory(self)->Optional[torch.Tensor]:	1✔
64	"""Return list of memories in which rows (outer dimension) are memories for each field.
65	These are derived from the matrix parameters of the afferent Projections to the retrieval_nodes
66	"""
67	num_fields = len(self.storage_node.afferents)	1✔
68	memory_matrices = [field.matrix for field in self.retrieve_projection_wrappers]	1✔
69	memory_capacity = len(memory_matrices[0])	1✔
70	return (None if not all(val for val in [num_fields, memory_matrices, memory_capacity])	1✔
71	else torch.stack([torch.stack([memory_matrices[j][i]
72	for j in range(num_fields)])
73	for i in range(memory_capacity)]))
74
75	def store_memory(self, memory_to_store, context):	1✔
76	"""Store variable in memory_matrix (parallel EMStorageMechanism._execute)
77
78	For each node in query_input_nodes and value_input_nodes,
79	assign its value to weights of corresponding afferents to corresponding match_node and/or retrieved_node.
80	- memory = matrix of entries made up vectors for each field in each entry (row)
81	- entry_to_store = query_input or value_input to store
82	- field_projections = Projections the matrices of which comprise memory
83
84	DIVISION OF LABOR between this method and function called by it
85	store_memory (corresponds to EMStorageMechanism._execute)
86	- compute norms to find weakest entry in memory
87	- compute storage_prob to determine whether to store current entry in memory
88	- call function with memory matrix for each field, to decay existing memory and assign input to weakest entry
89	storage_node.function (corresponds to EMStorage._function):
90	- decay existing memories
91	- assign input to weakest entry (given index for passed from EMStorageMechanism)
92
93	:return: List[2d tensor] updated memories
94	"""
95
96	memory = self.memory	1✔
97	assert memory is not None, f"PROGRAM ERROR: '{self.name}'.memory is None"	1✔
98
99	# Get current parameter values from EMComposition's EMStorageMechanism
100	mech = self.storage_node._mechanism	1✔
101	random_state = mech.function.parameters.random_state._get(context)	1✔
102	decay_rate = mech.parameters.decay_rate._get(context) # modulable, so use getter	1✔
103	storage_prob = mech.parameters.storage_prob._get(context) # modulable, so use getter	1✔
104	field_weights = mech.parameters.field_weights.get(context) # modulable, so use getter	1✔
105	concatenation_node = mech.concatenation_node	1✔
106	num_match_fields = 1 if concatenation_node else len([i for i in mech.field_types if i==1])	1✔
107
108	# Find weakest memory (i.e., with lowest norm)
109	field_norms = torch.empty((len(memory),len(memory[0])))	1✔
110	for row in range(len(memory)):	1✔
111	for col in range(len(memory[0])):	1✔
112	field_norms[row][col] = torch.linalg.norm(memory[row][col])	1✔
113	if field_weights is not None:	1!
UNCOV 114	field_norms *= field_weights	×
115	row_norms = torch.sum(field_norms, axis=1)	1✔
116	idx_of_weakest_memory = torch.argmin(row_norms)	1✔
117
118	values = []	1✔
119	for i, field_projection in enumerate(self.match_projection_wrappers + self.retrieve_projection_wrappers):	1✔
120	if i < num_match_fields:	1✔
121	# For match projections, get entry to store from value of sender of Projection matrix
122	# (this is to accomodate concatenation_node)
123	axis = 0	1✔
124	entry_to_store = field_projection.sender.value	1✔
125	if concatenation_node is None:	1!
126	assert (entry_to_store == memory_to_store[i]).all(), \	1✔
127	f"PROGRAM ERROR: misalignment between inputs and fields for storing them"
128	else:
129	# For retrieve projections, get entry to store from memory_to_store (which has inputs to all fields)
130	axis = 1	1✔
131	entry_to_store = memory_to_store[i - num_match_fields]	1✔
132	# Get matrix containing memories for the field from the Projection
133	field_memory_matrix = field_projection.matrix	1✔
134
135	field_projection.matrix = self.storage_node.function(entry_to_store,	1✔
136	memory_matrix=field_memory_matrix,
137	axis=axis,
138	storage_location=idx_of_weakest_memory,
139	storage_prob=storage_prob,
140	decay_rate=decay_rate,
141	random_state=random_state)
142	values.append(field_projection.matrix)	1✔
143
144	self.storage_node.value = values	1✔
145	return values	1✔

PrincetonUniversity / PsyNeuLink / 7186854566

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