Unable to predict AND gate output using a neural net

Pytorch beginner here. I am trying to predict the output of an AND gate using a neural network, but when I train the model, the loss doesn’t decrease and output stays at 0.2500 for every input. What am I doing wrong?

import torch 
import torch.nn as nn
import torch.nn.functional as F

class Net(nn.Module):
  def __init__(self):
    super().__init__()
    self.fc1 = nn.Linear(2, 4)
    self.fc2 = nn.Linear(4,4)
    self.fc3 = nn.Linear(4, 1)
    
  def forward(self, x):
    x = F.relu(self.fc1(x))
    x = F.relu(self.fc2(x))
    x = torch.sigmoid(self.fc3(x))
    
    return x

a = torch.tensor([[[0,0]],[[1,0]],[[0,1]],[[1,1]]], dtype = torch.float32)
y = torch.tensor([[0],[0],[0],[1]], dtype = torch.float32)

N = Net()

criterion = nn.MSELoss()
optimizer = torch.optim.Adam(N.parameters(), lr = 0.01)

for epoch in range(100):
  optimizer.zero_grad()
  out = N(a)
  loss = criterion(out, y)
  print(loss)
  loss.backward()
  optimizer.step()

Your code generally seems to work after fixing some minor shape issue in a (removing the unnecessary dim1).
I also tried to run your code with 10 random seeds and all seem to converge.

1 Like

@ptrblck
Thanks, it works now. But why was dim1 causing problems? Also I was unable to use the CrossEntropyLoss. It gave the error:

 RuntimeError: multi-target not supported at /pytorch/aten/src/THNN/generic/ClassNLLCriterion.c:21

Can’t I treat the first three inputs as of class 0 and the 4th input as of class 1?

Your input had the shape [4, 1, 2], which I changed to [4, 2], since dim1 is not needed.
If you want to use nn.CrossEntropyLoss, your target should have the shape [4] containing the class indices as values. Also, the sigmoid should be removed as the last layer in your model.
Additionally, if you would like to just use two classes, your last linear layer should also output two values (out_features=2).

1 Like

Thanks for removing my doubts.