Hello, I have very basic problem with training classification MLP network - I’m trying to train a network for simple classification task on randomly generated dataset with a bit imbalanced classes (59 observations of class 0 and 140 of class 1), and I can’t seem to teach the NN to distinguish between them, it always just simply predicts all the classes to class 1. I tried to use different weighting of classes, but it didn’t help. I know that the dataset is small, but unfortunately that was the task given by the teacher, and I shouldn’t make the dataset bigger. Here is my code:
class NeuralNet(nn.Module):
  def __init__(self, n_input, n_hidden1, n_hidden2, output_shape):
    super().__init__()
    self.hidden2 = nn.Linear(n_hidden1, n_hidden2)
    self.hidden = nn.Linear(n_input, n_hidden1)
    self.outer = nn.Linear(n_hidden1, output_shape)
    self.sigmoid = nn.Sigmoid()
    self.softmax = nn.Softmax(dim=1)
    
  def forward(self, data):
    step1 = self.hidden(data)
    step2 = self.sigmoid(step1)
    step3 = self.hidden2(step2)
    step4 = self.sigmoid(step3)
    step5 = self.outer(step4)
    step6 = self.softmax(step5)
    return step6

model = NeuralNet(4, 5, 5, 2)
model = model.float()
num_epoch = 100
learning_rate = 0.001
loss_function = F.binary_cross_entropy
optimizer = optim.SGD(model.parameters(), lr=learning_rate)#, weight_decay=1e-4, momentum=0.9)

#training loop:
train_scores = []
test_scores = []
for epoch in range(num_epoch):
  scores = []
  model.train()
  for x_batch, y_batch in train_dataloader:
    y_pred = model(x_batch).squeeze()
    y_batch = y_batch.squeeze().float()
    
    loss = loss_function(y_pred, y_batch, weight=torch.tensor([0.6, 0.4]))
    loss.backward()
    
    optimizer.step()
    optimizer.zero_grad()
    scores.append(loss.item())
    
  model.eval()
  train_loss = np.mean(scores)
  with torch.no_grad():
        test_loss = sum(loss_function(model(xb).squeeze(), yb.squeeze().float()) for xb, yb in test_dataloader)
  test_loss = test_loss/len(test_dataloader)
  #print(f"Epoch {epoch} train loss: {train_loss}, test loss: {test_loss}")
  train_scores.append(train_loss)
  test_scores.append(test_loss)