Partial Labeling

[BrokenSki8]

Hello,

I am trying to train a model for heart segmentation on multiple dataset at once. Those dataset have not been labeled with the same accuracy. For example, one dataset has only two labels, heart and background. While another has 13 labels, such as left ventricle, right ventricle, aorta, etc.

Is there anyways for me to force the output to give probability for each labels? Meaning I want to force the model to output the 13 labels for the first dataset and I'll handle the loss myself.

I have seen the Label Handling class but I am not sure if it would actually force the model to output logits for all of the labels, nor how to actually use it in my plan...

If anyone can point me to the code that might help, or explain how I could use the existing label handling class it would be great :)

Thanks in advance!