Upload raffael_conv_lstm.py with huggingface_hub

raffael_conv_lstm.py

@@ -0,0 +1,163 @@
+"""
+ConvLSTM Implementation
+True convolutional LSTM for spatiotemporal data processing
+"""
+import torch
+import torch.nn as nn
+
+
+class ConvLSTMCell(nn.Module):
+    """Single ConvLSTM Cell"""
+    
+    def __init__(self, input_dim, hidden_dim, kernel_size, bias=True):
+        super(ConvLSTMCell, self).__init__()
+        
+        self.input_dim = input_dim
+        self.hidden_dim = hidden_dim
+        self.kernel_size = kernel_size
+        self.padding = kernel_size[0] // 2, kernel_size[1] // 2
+        self.bias = bias
+        
+        # Input gate, forget gate, output gate, candidate values
+        self.conv = nn.Conv2d(
+            in_channels=self.input_dim + self.hidden_dim,
+            out_channels=4 * self.hidden_dim,
+            kernel_size=self.kernel_size,
+            padding=self.padding,
+            bias=self.bias
+        )
+    
+    def forward(self, input_tensor, cur_state):
+        h_cur, c_cur = cur_state
+        
+        # Concatenate input and hidden state
+        combined = torch.cat([input_tensor, h_cur], dim=1)
+        
+        # Compute all gates
+        combined_conv = self.conv(combined)
+        cc_i, cc_f, cc_o, cc_g = torch.split(combined_conv, self.hidden_dim, dim=1)
+        
+        # Activation functions
+        i = torch.sigmoid(cc_i)
+        f = torch.sigmoid(cc_f)
+        o = torch.sigmoid(cc_o)
+        g = torch.tanh(cc_g)
+        
+        # Update cell state and hidden state
+        c_next = f * c_cur + i * g
+        h_next = o * torch.tanh(c_next)
+        
+        return h_next, c_next
+    
+    def init_hidden(self, batch_size, image_size):
+        """Initialize hidden state"""
+        height, width = image_size
+        return (
+            torch.zeros(batch_size, self.hidden_dim, height, width, 
+                       device=self.conv.weight.device),
+            torch.zeros(batch_size, self.hidden_dim, height, width,
+                       device=self.conv.weight.device)
+        )
+
+
+class ConvLSTM(nn.Module):
+    """
+    ConvLSTM Module
+    Supports multiple layers, bidirectional (optional)
+    """
+    
+    def __init__(
+        self,
+        input_dim,
+        hidden_dim,
+        kernel_size,
+        num_layers=1,
+        batch_first=True,
+        bias=True,
+        return_all_layers=False
+    ):
+        super(ConvLSTM, self).__init__()
+        
+        self.input_dim = input_dim
+        # If hidden_dim is int, convert to list
+        if isinstance(hidden_dim, int):
+            self.hidden_dim = [hidden_dim] * num_layers
+        else:
+            self.hidden_dim = hidden_dim
+        self.kernel_size = kernel_size if isinstance(kernel_size, tuple) else (kernel_size, kernel_size)
+        self.num_layers = num_layers
+        self.batch_first = batch_first
+        self.bias = bias
+        self.return_all_layers = return_all_layers
+        
+        cell_list = []
+        for i in range(self.num_layers):
+            cur_input_dim = self.input_dim if i == 0 else self.hidden_dim[i - 1]
+            cell_list.append(
+                ConvLSTMCell(
+                    input_dim=cur_input_dim,
+                    hidden_dim=self.hidden_dim[i],
+                    kernel_size=self.kernel_size,
+                    bias=self.bias
+                )
+            )
+        self.cell_list = nn.ModuleList(cell_list)
+    
+    def forward(self, input_tensor, hidden_state=None):
+        """
+        Args:
+            input_tensor: (B, T, C, H, W) if batch_first else (T, B, C, H, W)
+            hidden_state: initial hidden state (optional)
+        
+        Returns:
+            last_state_list: (h_n, c_n) of last layer
+            layer_output_list: outputs of all timesteps
+        """
+        if not self.batch_first:
+            # (T, B, C, H, W) -> (B, T, C, H, W)
+            input_tensor = input_tensor.permute(1, 0, 2, 3, 4)
+        
+        b, _, _, h, w = input_tensor.size()
+        
+        # Initialize hidden state
+        if hidden_state is None:
+            hidden_state = self._init_hidden(batch_size=b, image_size=(h, w))
+        
+        layer_output_list = []
+        last_state_list = []
+        
+        seq_len = input_tensor.size(1)
+        cur_layer_input = input_tensor
+        
+        for layer_idx in range(self.num_layers):
+            h, c = hidden_state[layer_idx]
+            output_inner = []
+            
+            for t in range(seq_len):
+                h, c = self.cell_list[layer_idx](
+                    input_tensor=cur_layer_input[:, t, :, :, :],
+                    cur_state=[h, c]
+                )
+                output_inner.append(h)
+            
+            layer_output = torch.stack(output_inner, dim=1)  # (B, T, C, H, W)
+            cur_layer_input = layer_output
+            
+            layer_output_list.append(layer_output)
+            last_state_list.append([h, c])
+        
+        if not self.return_all_layers:
+            layer_output_list = layer_output_list[-1:]
+            last_state_list = last_state_list[-1:]
+        
+        return layer_output_list, last_state_list
+    
+    def _init_hidden(self, batch_size, image_size):
+        """Initialize hidden states for all layers"""
+        init_states = []
+        for i in range(self.num_layers):
+            init_states.append(
+                self.cell_list[i].init_hidden(batch_size, image_size)
+            )
+        return init_states
+