Weird Initial Predictions in 02. PyTorch NN Classification Model 0

[atharva-vadeyar]

Hey guys, for the first model in chapter 2, the model only predicts 1s. The logits are always greater than 0, so pushing them through the sigmoid then rounding them off always give me a 1 and an accuracy of 50.00% every time.

Plot decision boundary also doesn't give me any useful information.

Can someone please tell me what I'm missing here?

## Model Definition
class CircleModelV0(nn.Module):
  def __init__(self):
    super().__init__()
    self.layer1 = nn.Linear(in_features=2, out_features=5)
    self.layer2 = nn.Linear(in_features=5, out_features=1)

  def forward(self, x):
    return self.layer2(self.layer1(x))


model0 = CircleModelV0().to(device)


## Loss Function, Optimizer & Accuracy Function
# Create a loss function
loss_fn = nn.BCEWithLogitsLoss()

# Create an optimizer
optimizer = torch.optim.SGD(params=model0.parameters())

# Calculate accuracy:
def accuracy_fn(y_true, y_pred):
  correct = torch.eq(y_true, y_pred).sum().item()
  # print(f"Truth: {y_true}\nPred: {y_pred}")
  acc = (correct / len(y_pred)) * 100
  return acc


## Test/Train Loop
torch.manual_seed(42)

epochs = 100
X_train, y_train = X_train.to(device), y_train.to(device)
X_test, y_test = X_test.to(device), y_test.to(device)

for epoch in range(epochs):
  ### TRAINING
  model0.train()
  y_logits = model0(X_train).squeeze()
  # if (epoch == 1):
  #   print(y_logits)
  y_pred = torch.round(torch.sigmoid(y_logits))

  loss = loss_fn(y_logits, y_train)
  acc = accuracy_fn(y_true=y_train, y_pred=y_pred)

  optimizer.zero_grad()

  loss.backward()

  optimizer.step()

  ### TESTING
  model0.eval()
  with torch.inference_mode():
    test_logits = model0(X_test).squeeze()

    test_pred = torch.round(torch.sigmoid(test_logits))

    test_loss = loss_fn(test_logits, y_test)

    test_acc = accuracy_fn(y_true=y_test, y_pred=test_pred)

  if (epoch+1) % 10 == 0:
    print(f"Epoch: {epoch} | Loss: {loss:.5f}, Accuracy: {acc:.2f}, Test Loss: {test_loss:.5f}, Test Acc: {test_acc:.2f}")