Hi K. Frank,
Thanks for your explanation. I would like to validate my approach of getting the counts in the sample code above and some of my concerns for volumetric CT data.
import torch
torch.manual_seed(3)

image = torch.randint(low=0, high=2, size=(3, 3), dtype=torch.float32)
labelA = torch.randint(low=0, high=2, size=(3, 3), dtype=torch.float32)
labelB = torch.randint(low=0, high=2, size=(3, 3), dtype=torch.float32)

Generating the background label from the annotated label A & B i.e 0 in pixel value where either A or B is 1 and 1 elsewhere
background = torch.zeros((labelB.shape), dtype=torch.long)

for label in [labelA, labelB]:
    background |= background.long()

background = background^1
outputs = f"Image: {image},\n background: {background}, \n labelA: {labelA}, \n labelB: {labelB}"
print(outputs)

>>> outputs
Image: tensor([[0.0043, 0.1056, 0.2858],
        [0.0270, 0.4716, 0.0601],
        [0.7719, 0.7437, 0.5944]]),
 background: tensor([[0, 1, 0],
        [0, 0, 1],
        [0, 0, 0]]), 
 labelA: tensor([[1., 0., 1.],
        [1., 1., 0.],
        [1., 1., 0.]]), 
 labelB: tensor([[0., 0., 0.],
        [1., 1., 0.],
        [0., 0., 1.]])

counts = torch.tensor([torch.sum(background), torch.sum(labelA), torch.sum(labelB)])
weights = 1.0 / counts
weights_normalized = (1.0 / counts) / (1.0 / counts).sum()

>>> counts
tensor([4., 3., 4.])
>>> weights
tensor([0.2500, 0.3333, 0.2500])
>>>weights_normalized
tensor([0.3000, 0.4000, 0.3000])

These weights can now be passed into the loss function. My concerns are:

I believe my approach of creating the background is not well optimized, is there an inbuilt PyTorch function that can be used or another way around it?
Is it computationally efficient to create the background label class (for each corresponding image and label A/B) in the data loader or separately create it out the data loader, save it (.npy), and then reference it just like the label class A & B?

Thanks for your enlightenment