src/counting_small.py


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

import torch

def target_transform(labels):
    return labels

def loss_function(output, target):
    a = torch.tensor(range(10)).repeat(6 *
            target.shape[0]).reshape(target.shape[0], 6, 10)
    errors = (a - target.unsqueeze(2)).reshape(target.shape[0], 60)

    return torch.sum(output * errors ** 2) / output.shape[0]

def count_correct(output, target):
    output = output.reshape(output.shape[0], 6, 10)
    predictions = torch.argmax(output, dim=2)
    correct = torch.min(predictions.reshape(target.shape) == target, dim=1).values
    return torch.sum(correct).tolist()

def finalizer(x):
    x = x.reshape(x.shape[0], 6, 10)
    x = torch.softmax(x, 1)
    return x.reshape(x.shape[0], 60)

outputs = 60