test'

Author: Kevin Chang

bindsnet/models/models.py

@@ -549,150 +549,6 @@ def __init__(
         self.add_connection(recurrent_conn, source="Y", target="Y")
 
 
-import snntorch as snn
-
-class FFSNN(nn.Module):
-    # language=rst
-    """
-    A simple feedforward Spiking Neural Network (SNN) using snntorch,
-    designed for use with the ForwardForwardPipeline.
-    It consists of a sequence of Linear layers followed by Leaky Integrate-and-Fire
-    (LIF) spiking neuron layers.
-    """
-
-    def __init__(
-        self,
-        input_size: int, # This should be 794 (image_features + num_classes)
-        hidden_sizes: List[int], # e.g., [500, 500]
-        output_size: Optional[int] = None, # If the last FF layer is also the output layer for classification
-        beta: Union[float, torch.Tensor] = 0.9,  # Decay rate for snn.Leaky neurons
-        threshold: float = 1.0,  # Firing threshold for snn.Leaky neurons
-        reset_mechanism: str = "subtract",  # "subtract", "zero", or "none"
-        # Add other snn.Leaky parameters if needed, e.g., spike_grad
-    ) -> None:
-        # language=rst
-        """
-        Constructor for FFSNN.
-
-        :param input_size: Number of input features (after encoding and label embedding).
-        :param hidden_sizes: A list of integers, where each integer is the number of
-                             neurons in a hidden layer.
-        :param output_size: Optional. Number of neurons in the final layer if it's
-                            distinct or specifically for output. If None, the last
-                            size in hidden_sizes is considered the final FF layer.
-        :param beta: Membrane potential decay rate for Leaky neurons.
-        :param threshold: Firing threshold for Leaky neurons.
-        :param reset_mechanism: Reset mechanism for Leaky neurons after a spike.
-        """
-        super().__init__()
-
-        self.input_size = input_size
-        self.hidden_sizes = hidden_sizes
-        self.output_size = output_size
-        self.beta = beta
-        self.threshold = threshold
-        self.reset_mechanism = reset_mechanism
-
-        self.fc_layers = nn.ModuleList()
-        self.snn_layers = nn.ModuleList()
-        self._ff_layer_pairs_info = []
-
-        current_dim = self.input_size # Starts at 794
-        for i, hidden_dim in enumerate(self.hidden_sizes):
-            linear_layer = nn.Linear(current_dim, hidden_dim) # Layer 1: 794 -> 500
-                                                              # Layer 2: 500 -> 500
-            self.fc_layers.append(linear_layer)
-
-            snn_layer = snn.Leaky(
-                beta=self.beta, 
-                threshold=self.threshold, 
-                reset_mechanism=self.reset_mechanism,
-                # output_shape=[hidden_dim] # Optional: snntorch can infer this
-            )
-            self.snn_layers.append(snn_layer)
-            self._ff_layer_pairs_info.append((linear_layer, snn_layer))
-            current_dim = hidden_dim # Update current_dim for the *next* layer's input
-        
-        # If there's an output_size for a final classifier (not typical for pure FF layers)
-        if self.output_size is not None:
-            self.fc_out = nn.Linear(current_dim, self.output_size)
-            # Potentially another SNN layer if output is spiking
-            # self.snn_out = snn.Leaky(...)
-            # self._ff_layer_pairs_info.append((self.fc_out, self.snn_out)) # If FF applies here too
-
-    def forward(self, x_batch_time: torch.Tensor) -> Tuple[torch.Tensor, List[torch.Tensor]]:
-        # language=rst
-        """
-        Defines the forward pass of the SNN over time.
-        This method might be used if the network is called directly with time-series data.
-        However, the ForwardForwardPipeline._run_snn_batch currently iterates
-        through self.network_sequence modules per time step.
-
-        :param x_batch_time: Input tensor with shape [batch_size, time_steps, num_features].
-        :return: Final layer output spikes and a list of hidden states (membrane potentials)
-                 from all spiking layers.
-        """
-        # Initialize hidden states for all spiking layers in the sequence
-        # This assumes they are snn.Leaky and support init_leaky() or similar
-        # Or, more generally, that they initialize if mem is None on first call.
-        
-        spiking_layer_modules = [info['spiking'] for info in self._ff_layer_pairs_info]
-        # mem_states = [layer.init_leaky() for layer in spiking_layer_modules] # This creates new states
-        # For snntorch, typically pass None for initial state, layer handles it.
-        
-        # The pipeline's _run_snn_batch actually handles the time loop and state passing.
-        # This forward method is more for standalone use or if the pipeline changes.
-        # If you want this model to be directly callable with (B, T, F) input and manage its own time loop:
-
-        batch_size = x_batch_time.shape[0]
-        
-        # Initialize states for each spiking layer for this batch
-        # This is tricky because snn.Leaky.init_hidden() doesn't take batch_size.
-        # State initialization is usually handled by passing None to the layer's forward method
-        # for the first time step, and it initializes based on the input batch size.
-        
-        # Placeholder: The pipeline's _run_snn_batch is the primary runner.
-        # This forward method would need a more elaborate state management if used directly.
-        # For now, let's make it compatible with how _run_snn_batch works if it were to call this.
-        
-        # If this 'forward' is to be used, it should mirror _run_snn_batch's logic:
-        # Initialize all spiking layer states to None
-        spiking_layer_states = {module: None for module in self.network_sequence if isinstance(module, snn.SpikingNeuron)}
-        
-        # Record outputs if needed (e.g., for a final classification layer not part of FF)
-        # final_spk_rec = [] # If you want to record output spikes over time
-
-        for t in range(x_batch_time.shape[1]): # Iterate over time
-            x_t = x_batch_time[:, t, :]
-            layer_input = x_t
-            
-            current_module_idx = 0
-            for module in self.network_sequence:
-                if isinstance(module, snn.SpikingNeuron):
-                    spk_out, new_mem = module(layer_input, spiking_layer_states.get(module))
-                    spiking_layer_states[module] = new_mem
-                    layer_input = spk_out
-                else: # nn.Linear
-                    layer_input = module(layer_input)
-            # After passing through all layers for time step t, layer_input is the output of the last layer
-            # final_spk_rec.append(layer_input) 
-        
-        # return torch.stack(final_spk_rec, dim=1), [state for state in spiking_layer_states.values()]
-        return layer_input, [spiking_layer_states[info['spiking']] for info in self._ff_layer_pairs_info]
-
-
-    def get_ff_layer_pairs(self) -> List[Tuple[nn.Linear, snn.SpikingNeuron]]:
-        # language=rst
-        """
-        Returns the list of (Linear, SpikingNeuron) pairs for Forward-Forward training.
-        """
-        # If _ff_layer_pairs_info contains tuples, return them directly
-        return self._ff_layer_pairs_info
-        
-        # Alternative: If you want to be explicit about the structure
-        # return [(pair[0], pair[1]) for pair in self._ff_layer_pairs_info]
-
-
 class FFSNN_BindsNET(Network):
     # language=rst
     """