Angular Features

vision
8

Assuming my description is right, you could use this simple example to use the target for the current batch and index the output with it in order to append the features class-based:

nb_classes = 3
batch_size = 10
features = 4
output = torch.randn(batch_size, features)
targets = torch.randint(0, nb_classes, (batch_size,))

print(targets)
> tensor([0, 0, 0, 2, 0, 1, 1, 2, 2, 2])

class_features = {idx: [] for idx in range(nb_classes)}
for class_index in range(nb_classes):
    idx = targets == class_index
    class_feat = output[idx]
    class_features[class_index].extend(class_feat)

print(class_features)
> {0: [tensor([-0.8338,  0.3141, -0.2840,  0.6104]), tensor([-1.5458, -0.3546, -0.3190,  0.6153]), tensor([-1.3919, -0.4112,  0.4425,  0.8475]), tensor([ 0.3294, -0.2577,  0.3397,  0.4239])], 1: [tensor([-1.4398,  0.2516,  1.6932, -0.0364]), tensor([-0.6315, -0.6568,  0.7358,  0.5755])], 2: [tensor([ 1.3726, -0.7952,  2.1696,  0.6634]), tensor([ 0.4024, -2.2702,  2.2658,  2.6325]), tensor([ 0.2029, -1.4608, -0.2938,  1.0877]), tensor([-1.2883, -0.5849,  0.2535, -0.0638])]}

In the next step you could then use another loop to calculate the mean of these features:

class_features_mean = {idx: [] for idx in range(nb_classes)}
for class_index in range(nb_classes):
    tmp = torch.stack(class_features[class_index])
    class_features_mean[class_index] = tmp.mean(dim=0)

print(class_features_mean)
> {0: tensor([-0.8606, -0.1773,  0.0448,  0.6243]), 1: tensor([-1.0357, -0.2026,  1.2145,  0.2696]), 2: tensor([ 0.1724, -1.2778,  1.0988,  1.0799])}

Note that you could of course speed up this code, but I would recommend to stick to loops for now until you are sure the method works as intended.

1 Like