Update dygrae.py

torch_geometric_temporal/nn/recurrent/dygrae.py

@@ -1,7 +1,7 @@
 import torch
-from torch.nn import LSTM
+import torch.nn as nn
 from torch_geometric.nn import GatedGraphConv
-
+from torch.nn.utils import clip_grad_norm_
 
 class DyGrEncoder(torch.nn.Module):
     r"""An implementation of the integrated Gated Graph Convolution Long Short
@@ -42,12 +42,20 @@ def _create_layers(self):
             bias=True,
         )
 
-        self.recurrent_layer = LSTM(
+        self.recurrent_layer = nn.LSTMCell(
             input_size=self.conv_out_channels,
             hidden_size=self.lstm_out_channels,
-            num_layers=self.lstm_num_layers,
         )
 
+        self.dropout = nn.Dropout(0.5)  # Adjust dropout rate as needed
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        for param in self.parameters():
+            if param.dim() > 1:
+                nn.init.xavier_uniform_(param)
+
     def forward(
         self,
         X: torch.FloatTensor,
@@ -56,33 +64,44 @@ def forward(
         H: torch.FloatTensor = None,
         C: torch.FloatTensor = None,
     ) -> torch.FloatTensor:
-        """
-        Making a forward pass. If the hidden state and cell state matrices are
-        not present when the forward pass is called these are initialized with zeros.
-
-        Arg types:
-            * **X** *(PyTorch Float Tensor)* - Node features.
-            * **edge_index** *(PyTorch Long Tensor)* - Graph edge indices.
-            * **edge_weight** *(PyTorch Float Tensor, optional)* - Edge weight vector.
-            * **H** *(PyTorch Float Tensor, optional)* - Hidden state matrix for all nodes.
-            * **C** *(PyTorch Float Tensor, optional)* - Cell state matrix for all nodes.
-
-        Return types:
-            * **H_tilde** *(PyTorch Float Tensor)* - Output matrix for all nodes.
-            * **H** *(PyTorch Float Tensor)* - Hidden state matrix for all nodes.
-            * **C** *(PyTorch Float Tensor)* - Cell state matrix for all nodes.
-        """
         H_tilde = self.conv_layer(X, edge_index, edge_weight)
-        H_tilde = H_tilde[None, :, :]
+        H_tilde = self.dropout(H_tilde)
+
+        batch_size = H_tilde.size(0)
+
         if H is None and C is None:
-            H_tilde, (H, C) = self.recurrent_layer(H_tilde)
-        elif H is not None and C is not None:
-            H = H[None, :, :]
-            C = C[None, :, :]
-            H_tilde, (H, C) = self.recurrent_layer(H_tilde, (H, C))
-        else:
-            raise ValueError("Invalid hidden state and cell matrices.")
-        H_tilde = H_tilde.squeeze()
-        H = H.squeeze()
-        C = C.squeeze()
-        return H_tilde, H, C
+            H = torch.zeros(batch_size, self.lstm_out_channels).to(X.device)
+            C = torch.zeros(batch_size, self.lstm_out_channels).to(X.device)
+
+        H_out = []
+        C_out = []
+
+        for i in range(batch_size):
+            H_i, C_i = self.recurrent_layer(H_tilde[i], (H[i], C[i]))
+            H_out.append(H_i)
+            C_out.append(C_i)
+
+        H_out = torch.stack(H_out)
+        C_out = torch.stack(C_out)
+
+        return H_tilde, H_out, C_out
+
+# Example usage with data loader and GPU support
+model = DyGrEncoder(conv_out_channels, conv_num_layers, conv_aggr, lstm_out_channels, lstm_num_layers)
+model = model.to('cuda')  # Move model to GPU
+
+# Example data loader
+dataset = YourDataset(...)  # Replace with your own dataset
+data_loader = DataLoader(dataset, batch_size=32, shuffle=True)
+
+optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
+
+for epoch in range(num_epochs):
+    for batch_data in data_loader:
+        batch_data = batch_data.to('cuda')  # Move data to GPU
+        optimizer.zero_grad()
+        output, h_out, c_out = model(batch_data.X, batch_data.edge_index, batch_data.edge_weight)
+        loss = compute_loss(output, batch_data.y)  # Replace with your own loss computation
+        loss.backward()
+        clip_grad_norm_(model.parameters(), max_norm=1.0)  # Clip gradients to prevent explosion
+        optimizer.step()