I am training U-net with 3 channels, input data 40x40x40 set of voxels and using CELoss followed by ReLU activation. Also, there is only one convolution layer unlike the original U-net. For many iterations, my test loss evaluates to nan but train loss always gives a defined value. What could be the probable reason and how can this issue be sorted out?
I tried changing the activation function and also introducing the batch normalization layer after the convolution. But it didn't help.
Code for the loss evaluation part
model.eval()
test_loss, test_acc, b1,b2,s1,s2 = 0, 0, 0, 0, 0, 0
loss_fn: torch.nn.Module = nn.CrossEntropyLoss(ignore_index = 2, reduction='mean')
with torch.inference_mode():
# Loop through DataLoader batches
for batch, input_dataset in enumerate(dataloader):
Input = input_dataset[0]
Target = input_dataset[1]
prediction = model(Input.float())
loss = loss_fn(prediction, Target.long())
print("Loss",loss)
test_loss += loss.item()
predicted_label = prediction.argmax(dim=1)
s1 += (Target==1).sum().item()
b1 += (Target==0).sum().item()
acc, sb = f1_loss(predicted_label, Target)
s2 += sb[1]
b2 += sb[0]
test_acc += acc
test_loss = test_loss / len(dataloader)
print("DL",len(dataloader))
All output from the main function:
Loss tensor(nan)
Loss tensor(10.8893)
Loss tensor(12.2123)
DL 3
Epoch: 2 | train_loss: 111.9303 | train_acc: 0.9415 | test_loss: nan | test_acc: 0.8643
Loss tensor(nan)
Loss tensor(39.5226)
Loss tensor(17.3606)
DL 3
Epoch: 3 | train_loss: 25.7117 | train_acc: 0.7250 | test_loss: nan | test_acc: 0.9513