m-chrzan.xyz
aboutsummaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--src/counting_small.py12
1 files 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()