I am working on a multiclass segmentation problem where there are 8 different classes, including background class. I’ve implemented 3D UNet model from its original paper. (https://arxiv.org/pdf/1606.06650.pdf) You can see the architecture below. I am using Generalized Dice Loss. I got GDL code from SegLoss/dice_loss.py at master · JunMa11/SegLoss · GitHub. I am passing model output to Softmax loss function. However, I confused at some points. I need to pass dim parameter to Softmax function. I am passing 1 to Softmax function as dim parameter. Is it correct? Because I checked some models in GitHub and some of them used dim=2 parameter.
My actual problem is my loss is negative values. I couldn’t find where the problem is. I am separating my image into patches and eliminating the patches which consists of mostly zeros. Because my GPU resources are limited. I normalized my input data using torchio’s tio.ZNormalization function. However, my loss is still negative values. Could anyone help me?
----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv3d-1       [-1, 32, 64, 64, 64]             896
              ReLU-2       [-1, 32, 64, 64, 64]               0
            Conv3d-3       [-1, 64, 64, 64, 64]          55,360
              ReLU-4       [-1, 64, 64, 64, 64]               0
         MaxPool3d-5       [-1, 64, 32, 32, 32]               0
            Conv3d-6       [-1, 64, 32, 32, 32]         110,656
              ReLU-7       [-1, 64, 32, 32, 32]               0
            Conv3d-8      [-1, 128, 32, 32, 32]         221,312
              ReLU-9      [-1, 128, 32, 32, 32]               0
        MaxPool3d-10      [-1, 128, 16, 16, 16]               0
           Conv3d-11      [-1, 128, 16, 16, 16]         442,496
             ReLU-12      [-1, 128, 16, 16, 16]               0
           Conv3d-13      [-1, 256, 16, 16, 16]         884,992
             ReLU-14      [-1, 256, 16, 16, 16]               0
        MaxPool3d-15         [-1, 256, 8, 8, 8]               0
           Conv3d-16         [-1, 256, 8, 8, 8]       1,769,728
             ReLU-17         [-1, 256, 8, 8, 8]               0
           Conv3d-18         [-1, 512, 8, 8, 8]       3,539,456
             ReLU-19         [-1, 512, 8, 8, 8]               0
  ConvTranspose3d-20      [-1, 512, 16, 16, 16]       2,097,664
           Conv3d-21      [-1, 256, 16, 16, 16]       5,308,672
           Conv3d-22      [-1, 256, 16, 16, 16]       1,769,728
  ConvTranspose3d-23      [-1, 256, 32, 32, 32]         524,544
           Conv3d-24      [-1, 128, 32, 32, 32]       1,327,232
           Conv3d-25      [-1, 128, 32, 32, 32]         442,496
  ConvTranspose3d-26      [-1, 128, 64, 64, 64]         131,200
           Conv3d-27       [-1, 64, 64, 64, 64]         331,840
           Conv3d-28       [-1, 64, 64, 64, 64]         110,656
           Conv3d-29        [-1, 8, 64, 64, 64]             520
================================================================
Total params: 19,069,448
Trainable params: 19,069,448
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 1.00
Forward/backward pass size (MB): 1219.00
Params size (MB): 72.74
Estimated Total Size (MB): 1292.74
----------------------------------------------------------------
None

Training code (bs=1):
patch_indexes = create_patch_indexes(shape, patch_sizes)

for epoch in range(num_epochs):
    for i, (image, mask) in enumerate(train_loader):
        image = image.to(device) #[bs,x,y,z]
        image = image.view(batch_size_, 1, shape[0], shape[1], shape[2]) #[bs,c,x,y,z]
        
        mask = mask.to(device) #[x,y,z]
        mask = mask.view(batch_size_, 1, shape[0], shape[1], shape[2]) #[bs,1,x,y,z]
        
        
        for coors in patch_indexes:
            [sx, sy, sz] = coors[0]
            [ex, ey, ez] = coors[1]            
            patch_image = image[:, :, sx:ex, sy:ey, sz:ez]
            patch_mask = mask[:, :, sx:ex, sy:ey, sz:ez]
            
            unique, counts = torch.unique(mask[sx:ex, sy:ey, sz:ez], return_counts=True) 
            
            # if the pixel count of the label below the 70 pixel don't train with it.            
            if torch.any(counts[1:]>70):            
                outputs = model(patch_image.float()) #outputs [8, x, y, z]

                loss = criterion(outputs, patch_mask.long())
                print(loss.item())
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
        
        print (f'Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{n_total_steps}], Loss: {loss.item():.4f}')
        
        if (i+1) % 5 == 0:            
            torch.save(model.state_dict(), os.path.join(weight_path, "model.pth"))

print('Finished Training')

Loss results:
-0.040152035653591156
-3.7653479714324845e-31
Epoch [1/500], Step [1/70], Loss: -0.0000
-0.04126344621181488
-3.1088024030518945e-31
Epoch [1/500], Step [2/70], Loss: -0.0000
-0.03529515862464905
-4.84647575166059e-31
Epoch [1/500], Step [3/70], Loss: -0.0000
-0.008731698617339134
Epoch [1/500], Step [4/70], Loss: -0.0087
-0.01194444764405489
-5.359176546528797e-31
Epoch [1/500], Step [5/70], Loss: -0.0000
-0.034029535949230194
-1.801244036081391e-30
Epoch [1/500], Step [6/70], Loss: -0.0000
-0.01090750191360712
-1.0952236031865561e-31
Epoch [1/500], Step [7/70], Loss: -0.0000
-0.009786739945411682
Epoch [1/500], Step [8/70], Loss: -0.0098
-0.0001114458282245323
-3.798904103368798e-31
Epoch [1/500], Step [9/70], Loss: -0.0000
-0.0686672255396843
-1.1397233054103808e-30
Epoch [1/500], Step [10/70], Loss: -0.0000
-0.015976818278431892
-8.188111847262908e-31