Does the preprocessed input feature in forward process affect the model performance?

[songsong0425]

Dear PyG community,
Greetings, always thank you for your help.
I need your objective inspection of the model training process..

While I trained and validated the model, I recognized that the model always returned low evaluation scores (e.g., AUC, AUPR, F1-score) when I used real datasets or random values which means the model has not trained well.
The dataset is formed as below:

train_data
# Data(x=[2], edge_index=[2, 7854705], edge_label=[18429], edge_label_index=[2, 18429], edge_class=[18429], num_nodes=107940)
# train_data.x[0]: (107940, 256)
# train_data.x[1][i]: (19392, 4) # i in range(0, 5) and the shape of each dataset is same
# split the dataset based on the edge_label and edge_label_index

and the model is based on the simple GAT model.

class MyModel(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super(MyModel, self).__init__()
        
        # MLP
        self.change_dimension = nn.Sequential(
            nn.Linear(4, 64, bias=True),
            nn.LeakyReLU(), #inplace=True
            nn.Linear(64, 256, bias=True)
        )
        
        # GAT
        self.conv1 = GATConv(in_channels, hidden_channels, heads=6) # in_channel, hidden_channel, heads
        self.lin1 = torch.nn.Linear(in_channels, hidden_channels * 6)
        self.norm1 = LayerNorm(hidden_channels * 6)
        self.conv2 = GATConv(hidden_channels * 6, out_channels, heads=1, concat=False) # hidden_channel * heads, out_channel, heads
        self.lin2 = torch.nn.Linear(hidden_channels * 6, out_channels)
        

    def convert_data_size(self, data_t1, data_t2):
        x2 = torch.zeros((len(data_t1), 256), dtype=torch.float).to(device) #(107940, 256)
        x3 = self.change_dimension(data_t2).to(device) #(19392, 256)
        nonzero_index = torch.tensor(node_list).to(device) # The 19392 indices which contain nonzero values in x3
        x2.index_add_(0, nonzero_index, x3) #(107940, 256) # zero padding
        
        return x2
 
    def forward(self, data_t1, data_t2, node_ids, neighbor_cl_ids, edge_index):
        x1 = data_t1 #(107940, 256)
        
        for i in range(len(data_t2)):
            data_t2[i] = self.convert_data_size(data_t1, data_t2[i]) #(107940, 256)
         
        # The role of the below loop is just to fill the empty node feature matrix by mapping based on the edge_class
        # For example, if 0th node has positive label with 1th node in the data_t2[3], then put the node feature of 0th and 1st node from the data_t2[3]
        # I tried to make the model selectively get the corresponding node feature by edge_index_class per batch from the LinkNeighborLoader
        tmp_filled = torch.zeros((len(data_t1), 256), dtype=torch.float).to(device) #(107940, 256)
        for e_i_class in torch.unique(neighbor_cl_ids):
            indices = (neighbor_cl_ids == e_i_class).nonzero(as_tuple=True)[0]
            
            key_n_id = torch.unique(node_ids[edge_index[:, indices]])
            t2_n_id = np.in1d(node_list, key_n_id.cpu()).nonzero()[0]
            tmp_filled[node_list[t2_n_id]] = data_t2[int(e_i_class.item())][t2_n_id].float()

        x = torch.cat((x1, tmp_filled), dim=1) #(107940, 512)
        x_1, a1 = self.conv1(x, edge_index, return_attention_weights=True)
        x = F.leaky_relu(self.norm1(x_1)) 
        x = F.dropout(x, p = 0.2, training = self.training)
        x_2, a2 = self.conv2(x, edge_index, return_attention_weights=True)
        x = x_2 + self.lin2(x)

        return x, a1, a2

The training process is also simple. I just added the edge_class matching code.

def train(num_epochs):
    for epoch in range(1, num_epochs+1):
        tr_losses = 0
        val_losses = 0
        
        model.train()
        for data in tqdm(train_loader):
            data = data.to(device)
            data.edge_class = data.edge_class[data.input_id]
            data.edge_index_class = torch.zeros(len(data.edge_index[0])).to(device)
            # data: Data(x=[2], edge_index=[2, 453], edge_label=[64], edge_label_index=[2, 64], edge_class=[64], 
            # num_nodes=581, n_id=[581], e_id=[453], batch=[581], num_sampled_nodes=[3], num_sampled_edges=[2], 
            # input_id=[64], edge_index_class=[453])
            
            for i in range(len(data.edge_index_class)):
                if ((data.edge_index[1][i] in data.edge_label_index[0]) or (data.edge_index[1][i] in data.edge_label_index[1])):
                    data.edge_index_class[i] = data.edge_class[(data.edge_index[1][i]==data.edge_label_index[0])
                                                                 |(data.edge_index[1][i]==data.edge_label_index[1])]
                else:
                    data.edge_index_class[i] = torch.max(data.edge_index_class[(data.edge_index[1][i]==data.edge_index[0])])
                    
            optimizer.zero_grad()
            z, a1, a2 = model(data.x[0], data.x[1], data.n_id, data.edge_index_class, data.edge_index)
            out = ((z[data.edge_label_index[0]] * z[data.edge_label_index[1]]).sum(dim=-1)).view(-1) # product of a pair of nodes on each edge
            tr_loss = criterion(out, data.edge_label.float())
            tr_losses += tr_loss.item()
            tr_loss.backward()
            optimizer.step()
        avg_tr_loss = tr_losses/len(train_loader.dataset)
...

I guess that the for loop in the forward process evoked the problem since the simpler model without the for loop went better, but I'm not sure why the model couldn't train.
Does the modified node feature (in my code, tmp_filled) affect to the model training?
If anyone knows how to deal with this problem, please let me know.
Thank you for reading the long question.
Have a nice day!

[songsong0425]

I close this question since my approach is totally wrong.