Suppose we have the tensor looks this follow, assume that we have 3 different class.
sample = torch.Tensor([
[0.1, 0.1], #-> batch 1 / class 1
[0.2, 0.2], #-> batch 1 / class 2
[0.4, 0.4], #-> batch 1 / class 2
[0.5, 0.5] #-> batch 1 / class 0
[0.7, 0.7] #-> batch 1 / class 0
[0.3, 0.3] #-> batch 2 / class 1
[0.1, 0.1] #-> batch 2 / class 1
[0.8, 0.8] #-> batch 2 / class 0
])
batch_num = torch.Tensor([5, 3], dtype = torch.Long)
node_type = torch.Tensor([1, 2, 2, 0, 0, 1, 1, 0], dtype = torch.Long)
sample is our data.
batch_num means how many nodes in each sample in the batch.
node_type means type of each node, len(node_type) == sample .size(0)
What I want to compute is that:
for each sample, I want to compute the mean according to the class, so the final result would looks like
output= torch.Tensor([
[0.1, 0.1], #-> batch 1 / class 1
[0.3, 0.3], #-> batch 1 / class 2
[0.6, 0.6] #-> batch 1 / class 0
[0.2, 0.2] #-> batch 2 / class 1
[0.8, 0.8] #-> batch 2 / class 0
[0., 0.] #-> batch 2 / class 2
])
Note that, for batch 2, there only two classes, for this case, we need to add 0 tensor to it
I was stuck here for a while, any help would be appreciated.