Why does my model predict the same label? - graph

Question

I am training a graph convolution neural network to classify EEG signals into emotion classes. The input of my data is an array of size [12803216]-->[number of subjects * numbers of channels (nodes) * features of each node]. The output should be class 0(Negative) or class 1(Positive).The data is slightly imbalanced (45% class 0 and 55% class 1). The problem is that my model always predict label 0 as output for all inputs in the training stage regardless of the convolution function.
What is wrong with my code and how can I fix it? Any comments are welcome.
connectivity at the below code is predefined based at the connections of the 32 electrodes(nodes)
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import GCNConv, SAGEConv, ResGatedGraphConv, global_mean_pool, BatchNorm
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report, f1_score, accuracy_score
labels = np.load("/content/drive/MyDrive/ValenceLabels_thres_5.npy")
labels = np.array(labels, dtype='int64')
labels.shape
class EEGraph(nn.Module):
def __init__(self, embedding_dim, first_conv, n_layers, conv_layer):
super(EEGraph, self).__init__()
self.n_layers = n_layers
self.convs = []
self.bns = []
d_in = embedding_dim
d_out = first_conv
for i in range(n_layers):
self.convs.append(conv_layer(d_in, d_out))
self.bns.append(BatchNorm(d_out, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True))
if i < n_layers - 1:
d_in, d_out = d_out, 2*d_out
self.convs = torch.nn.ModuleList(self.convs)
self.bns = torch.nn.ModuleList(self.bns)
self.project = nn.Linear(d_out, 3) # d_in beacu
self.project.apply(lambda x: nn.init.xavier_normal_(x.weight, gain=1) if type(x) == nn.Linear else None)
def forward(self, x, edge_index):
for i, (conv, bn) in enumerate(zip(self.convs, self.bns)):
x = conv(x, edge_index).permute(0, 2, 1)
x = bn(x)
x = F.dropout(F.leaky_relu(x, negative_slope=0.01), p=0.5, training=self.training).permute(0, 2, 1)
out = x.mean(dim=1).squeeze(dim=-1)
out = self.project(out)
return F.softmax(out, dim=-1)
device = torch.device("cuda")
connectivity = [[channel_order.index(e[0]), channel_order.index(e[1])] for e in edges]
connectivity = torch.tensor(connectivity).t().contiguous().to(device)
best_f1_score = -1
best_trial_name = None
n_epochs = 500
lr = 1e-3
weight_decay = 1e-5
batch_size = 63
criterion = nn.CrossEntropyLoss()
for node_dim in [16]:
node_features = np.load(f"/content/deap_graph_valence{node_dim}_1.npy")
A, Xte, yA, yte = train_test_split(node_features, labels, test_size=0.2, shuffle=True, stratify=labels, random_state=0)
Xtr, Xtr_valid, ytr, ytr_valid = train_test_split(A, yA, test_size=0.2, shuffle=True, stratify=yA, random_state=0)
Xtr = torch.tensor(Xtr).float().to(device)
Xtr_valid = torch.tensor(Xtr_valid).float().to(device)
Xte = torch.tensor(Xte).float().to(device)
ytr = torch.tensor(ytr).to(device)
#ytr_valid = torch.tensor(ytr_valid).to(device)
#yte = torch.tensor(yte).to(device)
for conv_fn in [GCNConv, SAGEConv, ResGatedGraphConv]:
for n_layers in range(1, 4):
for conv_dim in [32, 64, 128,256]:
trial_name = f"node_dim_{node_dim}-conv_fn_{conv_fn.__name__}-conv_layers_{n_layers}-conv_dim_{conv_dim}"
print(f"#: {trial_name}")
model = EEGraph(embedding_dim=Xtr.shape[-1],
first_conv=conv_dim,
n_layers=n_layers,
conv_layer=conv_fn).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)
for epoch in range(n_epochs):
model.train()
indices = torch.randperm(len(Xtr))
for j, batch in enumerate(indices.view(-1, 63)):
optimizer.zero_grad()
batch_input = Xtr[batch]
outputs = model(batch_input, connectivity)
loss = criterion(outputs, ytr[batch])
loss.backward()
optimizer.step()
with torch.no_grad():
model.eval()
outputs = model(Xtr_valid, connectivity)
output_classes = torch.argmax(outputs, dim=-1).cpu().numpy()
f1 = f1_score(ytr_valid, output_classes, average="macro")
if f1 > best_f1_score:
best_trial_name = trial_name
best_f1_score = f1
print("-"*100)
print(f"Best model so far: {best_trial_name}")
print(f"Best F1 Score: %{100*best_f1_score:.2f}")
test_outputs = model(Xte, connectivity)
test_output_classes = torch.argmax(test_outputs, dim=-1).cpu().numpy()
print(classification_report(yte, test_output_classes, target_names=["Negative", "Positive"]))
print("-"*100)
print()