优化GWN， ddgcrn的代码，使其更简洁易读

model/DDGCRN/DDGCRN.py

@@ -1,248 +1,119 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
+import torch, torch.nn as nn, torch.nn.functional as F
 from collections import OrderedDict
 
 
 class DGCRM(nn.Module):
     def __init__(self, node_num, dim_in, dim_out, cheb_k, embed_dim, num_layers=1):
-        super(DGCRM, self).__init__()
-        assert num_layers >= 1, 'At least one DGCRM layer is required.'
-
-        self.node_num = node_num
-        self.input_dim = dim_in
-        self.num_layers = num_layers
-
-        # Initialize DGCRM cells
-        self.DGCRM_cells = nn.ModuleList([
-            DDGCRNCell(node_num, dim_in, dim_out, cheb_k, embed_dim)
-            if i == 0 else
-            DDGCRNCell(node_num, dim_out, dim_out, cheb_k, embed_dim)
-            for i in range(num_layers)
-        ])
+        super().__init__()
+        self.node_num, self.input_dim, self.num_layers = node_num, dim_in, num_layers
+        self.cells = nn.ModuleList(
+            [DDGCRNCell(node_num, dim_in if i == 0 else dim_out, dim_out, cheb_k, embed_dim) for i in range(num_layers)]
+        )
 
     def forward(self, x, init_state, node_embeddings):
-        """
-        Forward pass of the DGCRM model.
-
-        Parameters:
-        - x: Input tensor of shape (B, T, N, D)
-        - init_state: Initial hidden states of shape (num_layers, B, N, hidden_dim)
-        - node_embeddings: Node embeddings
-        """
         assert x.shape[2] == self.node_num and x.shape[3] == self.input_dim
-
-        seq_length = x.shape[1]
-        current_inputs = x
-        output_hidden = []
-
         for i in range(self.num_layers):
-            state = init_state[i]
-            inner_states = []
+            state, inner = init_state[i].to(x.device), []
+            for t in range(x.shape[1]):
+                state = self.cells[i](x[:, t, :, :], state, [node_embeddings[0][:, t, :, :], node_embeddings[1]])
+                inner.append(state)
+            init_state[i] = state
+            x = torch.stack(inner, dim=1)
+        return x, init_state
 
-            for t in range(seq_length):
-                state = self.DGCRM_cells[i](current_inputs[:, t, :, :], state,
-                                            [node_embeddings[0][:, t, :, :], node_embeddings[1]])
-                inner_states.append(state)
-
-            output_hidden.append(state)
-            current_inputs = torch.stack(inner_states, dim=1)
-
-        return current_inputs, output_hidden
-
-    def init_hidden(self, batch_size):
-        """
-        Initialize hidden states for DGCRM layers.
-
-        Parameters:
-        - batch_size: Size of the batch
-
-        Returns:
-        - Initial hidden states tensor
-        """
-        return torch.stack([
-            self.DGCRM_cells[i].init_hidden_state(batch_size)
-            for i in range(self.num_layers)
-        ], dim=0)
+    def init_hidden(self, bs):
+        return torch.stack([cell.init_hidden_state(bs) for cell in self.cells], dim=0)
 
 
 class DDGCRN(nn.Module):
     def __init__(self, args):
-        super(DDGCRN, self).__init__()
-
-        self.num_node = args['num_nodes']
-        self.input_dim = args['input_dim']
-        self.hidden_dim = args['rnn_units']
-        self.output_dim = args['output_dim']
-        self.horizon = args['horizon']
-        self.num_layers = args['num_layers']
-        self.use_day = args['use_day']
-        self.use_week = args['use_week']
-        self.default_graph = args['default_graph']
-
+        super().__init__()
+        self.num_node, self.input_dim, self.hidden_dim = args['num_nodes'], args['input_dim'], args['rnn_units']
+        self.output_dim, self.horizon, self.num_layers = args['output_dim'], args['horizon'], args['num_layers']
+        self.use_day, self.use_week = args['use_day'], args['use_week']
         self.node_embeddings1 = nn.Parameter(torch.randn(self.num_node, args['embed_dim']), requires_grad=True)
         self.node_embeddings2 = nn.Parameter(torch.randn(self.num_node, args['embed_dim']), requires_grad=True)
         self.T_i_D_emb = nn.Parameter(torch.empty(288, args['embed_dim']))
         self.D_i_W_emb = nn.Parameter(torch.empty(7, args['embed_dim']))
-
-        self.dropout1 = nn.Dropout(p=0.1)
-        self.dropout2 = nn.Dropout(p=0.1)
-
+        self.drop1, self.drop2 = nn.Dropout(0.1), nn.Dropout(0.1)
         self.encoder1 = DGCRM(self.num_node, self.input_dim, self.hidden_dim, args['cheb_order'], args['embed_dim'],
                               self.num_layers)
         self.encoder2 = DGCRM(self.num_node, self.input_dim, self.hidden_dim, args['cheb_order'], args['embed_dim'],
                               self.num_layers)
+        self.end_conv1 = nn.Conv2d(1, self.horizon * self.output_dim, (1, self.hidden_dim))
+        self.end_conv2 = nn.Conv2d(1, self.horizon * self.output_dim, (1, self.hidden_dim))
+        self.end_conv3 = nn.Conv2d(1, self.horizon * self.output_dim, (1, self.hidden_dim))
 
-        # Predictor
-        self.end_conv1 = nn.Conv2d(1, self.horizon * self.output_dim, kernel_size=(1, self.hidden_dim), bias=True)
-        self.end_conv2 = nn.Conv2d(1, self.horizon * self.output_dim, kernel_size=(1, self.hidden_dim), bias=True)
-        self.end_conv3 = nn.Conv2d(1, self.horizon * self.output_dim, kernel_size=(1, self.hidden_dim), bias=True)
-
-    def forward(self, source, **kwargs):
-        """
-        Forward pass of the DDGCRN model.
-
-        Parameters:
-        - source: Input tensor of shape (B, T_1, N, D)
-        - mode: Control mode for the forward pass
-
-        Returns:
-        - Output tensor
-        """
-        node_embedding1 = self.node_embeddings1
-
+    def forward(self, source):
+        node_embed = self.node_embeddings1
         if self.use_day:
-            t_i_d_data = source[..., 1]
-            T_i_D_emb = self.T_i_D_emb[(t_i_d_data * 288).long()]
-            node_embedding1 = node_embedding1 * T_i_D_emb
-
+            node_embed = node_embed * self.T_i_D_emb[(source[..., 1] * 288).long()]
         if self.use_week:
-            d_i_w_data = source[..., 2]
-            D_i_W_emb = self.D_i_W_emb[d_i_w_data.long()]
-            node_embedding1 = node_embedding1 * D_i_W_emb
-
-        node_embeddings = [node_embedding1, self.node_embeddings1]
+            node_embed = node_embed * self.D_i_W_emb[source[..., 2].long()]
+        node_embeddings = [node_embed, self.node_embeddings1]
         source = source[..., 0].unsqueeze(-1)
-
-        init_state1 = self.encoder1.init_hidden(source.shape[0])
-        output, _ = self.encoder1(source, init_state1, node_embeddings)
-        output = self.dropout1(output[:, -1:, :, :])
-        output1 = self.end_conv1(output)
-
-        source1 = self.end_conv2(output)
-        source2 = source[:, -self.horizon:, ...] - source1
-
-        init_state2 = self.encoder2.init_hidden(source2.shape[0])
-        output2, _ = self.encoder2(source2, init_state2, node_embeddings)
-        output2 = self.dropout2(output2[:, -1:, :, :])
-        output2 = self.end_conv3(output2)
-
-        return output1 + output2
+        init1 = self.encoder1.init_hidden(source.shape[0])
+        out, _ = self.encoder1(source, init1, node_embeddings)
+        out = self.drop1(out[:, -1:, :, :])
+        out1 = self.end_conv1(out)
+        src1 = self.end_conv2(out)
+        src2 = source[:, -self.horizon:, ...] - src1
+        init2 = self.encoder2.init_hidden(source.shape[0])
+        out2, _ = self.encoder2(src2, init2, node_embeddings)
+        out2 = self.drop2(out2[:, -1:, :, :])
+        return out1 + self.end_conv3(out2)
 
 
 class DDGCRNCell(nn.Module):
     def __init__(self, node_num, dim_in, dim_out, cheb_k, embed_dim):
-        super(DDGCRNCell, self).__init__()
-        self.node_num = node_num
-        self.hidden_dim = dim_out
-        self.gate = DGCN(dim_in + self.hidden_dim, 2 * dim_out, cheb_k, embed_dim)
-        self.update = DGCN(dim_in + self.hidden_dim, dim_out, cheb_k, embed_dim)
+        super().__init__()
+        self.node_num, self.hidden_dim = node_num, dim_out
+        self.gate = DGCN(dim_in + dim_out, 2 * dim_out, cheb_k, embed_dim, node_num)
+        self.update = DGCN(dim_in + dim_out, dim_out, cheb_k, embed_dim, node_num)
 
     def forward(self, x, state, node_embeddings):
-        state = state.to(x.device)
-        input_and_state = torch.cat((x, state), dim=-1)
-        z_r = torch.sigmoid(self.gate(input_and_state, node_embeddings))
-        z, r = torch.split(z_r, self.hidden_dim, dim=-1)
-        candidate = torch.cat((x, z * state), dim=-1)
-        hc = torch.tanh(self.update(candidate, node_embeddings))
-        h = r * state + (1 - r) * hc
-        return h
+        inp = torch.cat((x, state), -1)
+        z_r = torch.sigmoid(self.gate(inp, node_embeddings))
+        z, r = torch.split(z_r, self.hidden_dim, -1)
+        hc = torch.tanh(self.update(torch.cat((x, z * state), -1), node_embeddings))
+        return r * state + (1 - r) * hc
 
-    def init_hidden_state(self, batch_size):
-        return torch.zeros(batch_size, self.node_num, self.hidden_dim)
+    def init_hidden_state(self, bs):
+        return torch.zeros(bs, self.node_num, self.hidden_dim)
 
 
 class DGCN(nn.Module):
-    def __init__(self, dim_in, dim_out, cheb_k, embed_dim):
-        super(DGCN, self).__init__()
-        self.cheb_k = cheb_k
-        self.embed_dim = embed_dim
-
-        # Initialize parameters
+    def __init__(self, dim_in, dim_out, cheb_k, embed_dim, num_nodes):
+        super().__init__()
+        self.cheb_k, self.embed_dim = cheb_k, embed_dim
         self.weights_pool = nn.Parameter(torch.FloatTensor(embed_dim, cheb_k, dim_in, dim_out))
         self.weights = nn.Parameter(torch.FloatTensor(cheb_k, dim_in, dim_out))
         self.bias_pool = nn.Parameter(torch.FloatTensor(embed_dim, dim_out))
         self.bias = nn.Parameter(torch.FloatTensor(dim_out))
-
-        # Hyperparameters
-        self.hyperGNN_dim = 16
-        self.middle_dim = 2
-
-        # Fully connected layers
         self.fc = nn.Sequential(OrderedDict([
-            ('fc1', nn.Linear(dim_in, self.hyperGNN_dim)),
+            ('fc1', nn.Linear(dim_in, 16)),
             ('sigmoid1', nn.Sigmoid()),
-            ('fc2', nn.Linear(self.hyperGNN_dim, self.middle_dim)),
+            ('fc2', nn.Linear(16, 2)),
             ('sigmoid2', nn.Sigmoid()),
-            ('fc3', nn.Linear(self.middle_dim, self.embed_dim))
+            ('fc3', nn.Linear(2, embed_dim))
         ]))
+        # 预注册恒定不变的单位矩阵
+        self.register_buffer('eye', torch.eye(num_nodes))
 
     def forward(self, x, node_embeddings):
-        """
-        Forward pass for the DGCN model.
-
-        Parameters:
-        - x: Input tensor of shape [B, N, C]
-        - node_embeddings: Node embeddings tensor of shape [N, D]
-        - connMtx: Connectivity matrix
-
-        Returns:
-        - x_gconv: Output tensor of shape [B, N, dim_out]
-        """
-
-        node_num = node_embeddings[0].shape[1]
-        supports1 = torch.eye(node_num).to(node_embeddings[0].device)  # Identity matrix
-
-        # Apply fully connected layers
-        filter = self.fc(x)
-        nodevec = torch.tanh(torch.mul(node_embeddings[0], filter))  # Element-wise multiplication
-
-        # Compute Laplacian
-        supports2 = self.get_laplacian(F.relu(torch.matmul(nodevec, nodevec.transpose(2, 1))), supports1)
-
-        # Graph convolution
-        x_g1 = torch.einsum("nm,bmc->bnc", supports1, x)
-        x_g2 = torch.einsum("bnm,bmc->bnc", supports2, x)
-        x_g = torch.stack([x_g1, x_g2], dim=1)
-
-        # Apply graph convolution weights and biases
+        supp1 = self.eye.to(node_embeddings[0].device)
+        filt = self.fc(x)
+        nodevec = torch.tanh(node_embeddings[0] * filt)
+        supp2 = self.get_laplacian(F.relu(torch.matmul(nodevec, nodevec.transpose(2, 1))), supp1)
+        x_g = torch.stack([torch.einsum("nm,bmc->bnc", supp1, x),
+                           torch.einsum("bnm,bmc->bnc", supp2, x)], dim=1)
         weights = torch.einsum('nd,dkio->nkio', node_embeddings[1], self.weights_pool)
         bias = torch.matmul(node_embeddings[1], self.bias_pool)
-
-        x_g = x_g.permute(0, 2, 1, 3)  # Rearrange dimensions
-        x_gconv = torch.einsum('bnki,nkio->bno', x_g, weights) + bias  # Graph convolution operation
-
-        return x_gconv
+        return torch.einsum('bnki,nkio->bno', x_g.permute(0, 2, 1, 3), weights) + bias
 
     @staticmethod
     def get_laplacian(graph, I, normalize=True):
-        """
-        Compute the Laplacian of the graph.
-
-        Parameters:
-        - graph: Adjacency matrix of the graph, [N, N]
-        - I: Identity matrix
-        - normalize: Whether to use the normalized Laplacian
-
-        Returns:
-        - L: Graph Laplacian
-        """
-        if normalize:
-            D_inv_sqrt = torch.diag_embed(torch.sum(graph, dim=-1) ** (-1 / 2))
-            L = torch.matmul(torch.matmul(D_inv_sqrt, graph), D_inv_sqrt)
-        else:
-            graph = graph + I
-            D_inv_sqrt = torch.diag_embed(torch.sum(graph, dim=-1) ** (-1 / 2))
-            L = torch.matmul(torch.matmul(D_inv_sqrt, graph), D_inv_sqrt)
-        return L
-
+        D_inv = torch.diag_embed(torch.sum(graph, -1) ** (-0.5))
+        return torch.matmul(torch.matmul(D_inv, graph), D_inv) if normalize else torch.matmul(
+            torch.matmul(D_inv, graph + I), D_inv)
+ 

model/DDGCRN/DDGCRN_old.py

@@ -0,0 +1,248 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from collections import OrderedDict
+
+
+class DGCRM(nn.Module):
+    def __init__(self, node_num, dim_in, dim_out, cheb_k, embed_dim, num_layers=1):
+        super(DGCRM, self).__init__()
+        assert num_layers >= 1, 'At least one DGCRM layer is required.'
+
+        self.node_num = node_num
+        self.input_dim = dim_in
+        self.num_layers = num_layers
+
+        # Initialize DGCRM cells
+        self.DGCRM_cells = nn.ModuleList([
+            DDGCRNCell(node_num, dim_in, dim_out, cheb_k, embed_dim)
+            if i == 0 else
+            DDGCRNCell(node_num, dim_out, dim_out, cheb_k, embed_dim)
+            for i in range(num_layers)
+        ])
+
+    def forward(self, x, init_state, node_embeddings):
+        """
+        Forward pass of the DGCRM model.
+
+        Parameters:
+        - x: Input tensor of shape (B, T, N, D)
+        - init_state: Initial hidden states of shape (num_layers, B, N, hidden_dim)
+        - node_embeddings: Node embeddings
+        """
+        assert x.shape[2] == self.node_num and x.shape[3] == self.input_dim
+
+        seq_length = x.shape[1]
+        current_inputs = x
+        output_hidden = []
+
+        for i in range(self.num_layers):
+            state = init_state[i]
+            inner_states = []
+
+            for t in range(seq_length):
+                state = self.DGCRM_cells[i](current_inputs[:, t, :, :], state,
+                                            [node_embeddings[0][:, t, :, :], node_embeddings[1]])
+                inner_states.append(state)
+
+            output_hidden.append(state)
+            current_inputs = torch.stack(inner_states, dim=1)
+
+        return current_inputs, output_hidden
+
+    def init_hidden(self, batch_size):
+        """
+        Initialize hidden states for DGCRM layers.
+
+        Parameters:
+        - batch_size: Size of the batch
+
+        Returns:
+        - Initial hidden states tensor
+        """
+        return torch.stack([
+            self.DGCRM_cells[i].init_hidden_state(batch_size)
+            for i in range(self.num_layers)
+        ], dim=0)
+
+
+class DDGCRN(nn.Module):
+    def __init__(self, args):
+        super(DDGCRN, self).__init__()
+
+        self.num_node = args['num_nodes']
+        self.input_dim = args['input_dim']
+        self.hidden_dim = args['rnn_units']
+        self.output_dim = args['output_dim']
+        self.horizon = args['horizon']
+        self.num_layers = args['num_layers']
+        self.use_day = args['use_day']
+        self.use_week = args['use_week']
+        self.default_graph = args['default_graph']
+
+        self.node_embeddings1 = nn.Parameter(torch.randn(self.num_node, args['embed_dim']), requires_grad=True)
+        self.node_embeddings2 = nn.Parameter(torch.randn(self.num_node, args['embed_dim']), requires_grad=True)
+        self.T_i_D_emb = nn.Parameter(torch.empty(288, args['embed_dim']))
+        self.D_i_W_emb = nn.Parameter(torch.empty(7, args['embed_dim']))
+
+        self.dropout1 = nn.Dropout(p=0.1)
+        self.dropout2 = nn.Dropout(p=0.1)
+
+        self.encoder1 = DGCRM(self.num_node, self.input_dim, self.hidden_dim, args['cheb_order'], args['embed_dim'],
+                              self.num_layers)
+        self.encoder2 = DGCRM(self.num_node, self.input_dim, self.hidden_dim, args['cheb_order'], args['embed_dim'],
+                              self.num_layers)
+
+        # Predictor
+        self.end_conv1 = nn.Conv2d(1, self.horizon * self.output_dim, kernel_size=(1, self.hidden_dim), bias=True)
+        self.end_conv2 = nn.Conv2d(1, self.horizon * self.output_dim, kernel_size=(1, self.hidden_dim), bias=True)
+        self.end_conv3 = nn.Conv2d(1, self.horizon * self.output_dim, kernel_size=(1, self.hidden_dim), bias=True)
+
+    def forward(self, source, **kwargs):
+        """
+        Forward pass of the DDGCRN model.
+
+        Parameters:
+        - source: Input tensor of shape (B, T_1, N, D)
+        - mode: Control mode for the forward pass
+
+        Returns:
+        - Output tensor
+        """
+        node_embedding1 = self.node_embeddings1
+
+        if self.use_day:
+            t_i_d_data = source[..., 1]
+            T_i_D_emb = self.T_i_D_emb[(t_i_d_data * 288).long()]
+            node_embedding1 = node_embedding1 * T_i_D_emb
+
+        if self.use_week:
+            d_i_w_data = source[..., 2]
+            D_i_W_emb = self.D_i_W_emb[d_i_w_data.long()]
+            node_embedding1 = node_embedding1 * D_i_W_emb
+
+        node_embeddings = [node_embedding1, self.node_embeddings1]
+        source = source[..., 0].unsqueeze(-1)
+
+        init_state1 = self.encoder1.init_hidden(source.shape[0])
+        output, _ = self.encoder1(source, init_state1, node_embeddings)
+        output = self.dropout1(output[:, -1:, :, :])
+        output1 = self.end_conv1(output)
+
+        source1 = self.end_conv2(output)
+        source2 = source[:, -self.horizon:, ...] - source1
+
+        init_state2 = self.encoder2.init_hidden(source2.shape[0])
+        output2, _ = self.encoder2(source2, init_state2, node_embeddings)
+        output2 = self.dropout2(output2[:, -1:, :, :])
+        output2 = self.end_conv3(output2)
+
+        return output1 + output2
+
+
+class DDGCRNCell(nn.Module):
+    def __init__(self, node_num, dim_in, dim_out, cheb_k, embed_dim):
+        super(DDGCRNCell, self).__init__()
+        self.node_num = node_num
+        self.hidden_dim = dim_out
+        self.gate = DGCN(dim_in + self.hidden_dim, 2 * dim_out, cheb_k, embed_dim)
+        self.update = DGCN(dim_in + self.hidden_dim, dim_out, cheb_k, embed_dim)
+
+    def forward(self, x, state, node_embeddings):
+        state = state.to(x.device)
+        input_and_state = torch.cat((x, state), dim=-1)
+        z_r = torch.sigmoid(self.gate(input_and_state, node_embeddings))
+        z, r = torch.split(z_r, self.hidden_dim, dim=-1)
+        candidate = torch.cat((x, z * state), dim=-1)
+        hc = torch.tanh(self.update(candidate, node_embeddings))
+        h = r * state + (1 - r) * hc
+        return h
+
+    def init_hidden_state(self, batch_size):
+        return torch.zeros(batch_size, self.node_num, self.hidden_dim)
+
+
+class DGCN(nn.Module):
+    def __init__(self, dim_in, dim_out, cheb_k, embed_dim):
+        super(DGCN, self).__init__()
+        self.cheb_k = cheb_k
+        self.embed_dim = embed_dim
+
+        # Initialize parameters
+        self.weights_pool = nn.Parameter(torch.FloatTensor(embed_dim, cheb_k, dim_in, dim_out))
+        self.weights = nn.Parameter(torch.FloatTensor(cheb_k, dim_in, dim_out))
+        self.bias_pool = nn.Parameter(torch.FloatTensor(embed_dim, dim_out))
+        self.bias = nn.Parameter(torch.FloatTensor(dim_out))
+
+        # Hyperparameters
+        self.hyperGNN_dim = 16
+        self.middle_dim = 2
+
+        # Fully connected layers
+        self.fc = nn.Sequential(OrderedDict([
+            ('fc1', nn.Linear(dim_in, self.hyperGNN_dim)),
+            ('sigmoid1', nn.Sigmoid()),
+            ('fc2', nn.Linear(self.hyperGNN_dim, self.middle_dim)),
+            ('sigmoid2', nn.Sigmoid()),
+            ('fc3', nn.Linear(self.middle_dim, self.embed_dim))
+        ]))
+
+    def forward(self, x, node_embeddings):
+        """
+        Forward pass for the DGCN model.
+
+        Parameters:
+        - x: Input tensor of shape [B, N, C]
+        - node_embeddings: Node embeddings tensor of shape [N, D]
+        - connMtx: Connectivity matrix
+
+        Returns:
+        - x_gconv: Output tensor of shape [B, N, dim_out]
+        """
+
+        node_num = node_embeddings[0].shape[1]
+        supports1 = torch.eye(node_num).to(node_embeddings[0].device)  # Identity matrix
+
+        # Apply fully connected layers
+        filter = self.fc(x)
+        nodevec = torch.tanh(torch.mul(node_embeddings[0], filter))  # Element-wise multiplication
+
+        # Compute Laplacian
+        supports2 = self.get_laplacian(F.relu(torch.matmul(nodevec, nodevec.transpose(2, 1))), supports1)
+
+        # Graph convolution
+        x_g1 = torch.einsum("nm,bmc->bnc", supports1, x)
+        x_g2 = torch.einsum("bnm,bmc->bnc", supports2, x)
+        x_g = torch.stack([x_g1, x_g2], dim=1)
+
+        # Apply graph convolution weights and biases
+        weights = torch.einsum('nd,dkio->nkio', node_embeddings[1], self.weights_pool)
+        bias = torch.matmul(node_embeddings[1], self.bias_pool)
+
+        x_g = x_g.permute(0, 2, 1, 3)  # Rearrange dimensions
+        x_gconv = torch.einsum('bnki,nkio->bno', x_g, weights) + bias  # Graph convolution operation
+
+        return x_gconv
+
+    @staticmethod
+    def get_laplacian(graph, I, normalize=True):
+        """
+        Compute the Laplacian of the graph.
+
+        Parameters:
+        - graph: Adjacency matrix of the graph, [N, N]
+        - I: Identity matrix
+        - normalize: Whether to use the normalized Laplacian
+
+        Returns:
+        - L: Graph Laplacian
+        """
+        if normalize:
+            D_inv_sqrt = torch.diag_embed(torch.sum(graph, dim=-1) ** (-1 / 2))
+            L = torch.matmul(torch.matmul(D_inv_sqrt, graph), D_inv_sqrt)
+        else:
+            graph = graph + I
+            D_inv_sqrt = torch.diag_embed(torch.sum(graph, dim=-1) ** (-1 / 2))
+            L = torch.matmul(torch.matmul(D_inv_sqrt, graph), D_inv_sqrt)
+        return L
+

model/model_selector.py

@@ -13,6 +13,7 @@ from model.STFGNN.STFGNN import STFGNN
 from model.STSGCN.STSGCN import STSGCN
 from model.STGODE.STGODE import ODEGCN
 from model.PDG2SEQ.PDG2Seq import PDG2Seq
+from model.EXP.EXP import EXP
 
 def model_selector(model):
     match model['type']:
@@ -31,4 +32,5 @@ def model_selector(model):
         case 'STSGCN': return STSGCN(model)
         case 'STGODE': return ODEGCN(model)
         case 'PDG2SEQ': return PDG2Seq(model)
+        case 'EXP': return EXP(model)