Merging outputs of semantic segmentation model

HI guys,

In the context of semantic segmentation (19 classes) my model output a tensor of size :
(b, 19, h, w).

I need to merge some classes to get a tensor of size (b, 10, h, w) with respect to a mapping dictionary:

dic = {9:[0,1], 8:[3,4], 7:[9,10,11], 6:[6,7,8], 5:[2,5], 4:[18,17,16], 3:[4,5,6], 2:[12,13], 1:[15], 0:[14]}

What is the most efficient way to do that ?

Thanks !

I’m not sure, if you could use an alternative to a simple for loop :confused:

b, h, w = 2, 4, 4
x = torch.randn(b, 19, h, w)
z = torch.zeros(b, 10, h, w)

dic = {9:[0,1], 8:[3,4], 7:[9,10,11], 6:[6,7,8], 5:[2,5], 4:[18,17,16], 3:[4,5,6], 2:[12,13], 1:[15], 0:[14]}

for key in dic:
    d = dic[key]
    z[:, key] = x[:, d].sum()
1 Like

Hi !

Thank you for your answer, yes I don’t see any other alternative I will go for the for loop then.

Thanks