Hi,
Thanks for the tip. I’ll look into the stitching and see if that’s the problem. In the meantime, the weighted DICE and weighted cross entropy I implemented in the following manner:
Because it’s a highly unbalanced dataset (21 organs, some of them much smaller than others), we calculate the weights online, by passing the ground truth to the following function:
    def compute_weights(self, y):
        w = np.zeros(self.num_classes)
        labels, histo = np.unique(y.numpy(), return_counts=True)

        #index=0
        #for x in np.nditer(labels):
         #   w[x]= histo[index]
         #   index+=1
        # for i in range(len(labels)):
        #     if i in labels:
        #         w[labels[i]] = histo[index]
        #         index+=1
        # w[w==0]= eps
        freq = histo / np.sum(histo)
        med_freq = np.median(freq)
        weights = med_freq / (freq+eps)
        for i, label in enumerate(weights):
            w[int(labels[i])] = weights[i]

        return torch.from_numpy(w.astype(float))

When backpropagating loss, I figured that it doesn’t make much sense to backpropagate error for classes that don’t exist in the given patch, so we start with all weights equaling zero, and only change those weights in the following manner (median frequency balacning) for those classes existing in the current patch.
Then one can give these weights to torch.nn.CrossEntropyLoss to calculate the weighted CE:
        self.cross_entropy_loss = nn.CrossEntropyLoss(weight=weight)
        CE_Loss = self.cross_entropy_loss.forward(input, target)

The dice loss we calculate via the dice score, and then scale the loss for each channel by the online weights for the given patch : weights* (1 - Dice_Score).
Before calculating these median freq weights, my network wouldn’t learn from precalculated weights on the entire volume. This is because a patch was not representative of the volume, and so we were getting a lot of error backpropagation for classes that weren’t even represented in the current patch.
Hope this makes sense! If you have any feedback, or I missed something, lemme know!