From 18ca053adb132b90d9fe8cf37685d4249ab4e27c Mon Sep 17 00:00:00 2001 From: Marcin Chrzanowski Date: Mon, 3 May 2021 13:11:37 +0200 Subject: Fix counting loss --- src/counting_small.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/src/counting_small.py b/src/counting_small.py index 06cd690..62dafa3 100644 --- a/src/counting_small.py +++ b/src/counting_small.py @@ -1,21 +1,17 @@ import torch def target_transform(labels): - count_labels = torch.zeros(60) - for shape_index in range(len(labels)): - count_labels[int(shape_index * 10 + labels[shape_index])] = 1 - return count_labels + return labels def loss_function(output, target): - a = torch.tensor(range(10)).repeat(6 * target.shape[0]).reshape(target.shape) - errors = a - 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) - target = target.reshape(target.shape[0], 6, 10) - target = torch.argmax(target, dim=2) predictions = torch.argmax(output, dim=2) correct = torch.min(predictions.reshape(target.shape) == target, dim=1).values return torch.sum(correct).tolist() -- cgit v1.2.3