Skip to content

Commit 35c9740

Browse files
committed
Fix accuracy when topk > num_classes
1 parent a16a753 commit 35c9740

File tree

1 file changed

+6
-2
lines changed

1 file changed

+6
-2
lines changed

timm/utils/metrics.py

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
33
Hacked together by / Copyright 2020 Ross Wightman
44
"""
5+
import torch
56

67

78
class AverageMeter:
@@ -24,9 +25,12 @@ def update(self, val, n=1):
2425

2526
def accuracy(output, target, topk=(1,)):
2627
"""Computes the accuracy over the k top predictions for the specified values of k"""
27-
maxk = max(topk)
28+
maxk = min(max(topk), output.size()[1])
2829
batch_size = target.size(0)
2930
_, pred = output.topk(maxk, 1, True, True)
3031
pred = pred.t()
3132
correct = pred.eq(target.reshape(1, -1).expand_as(pred))
32-
return [correct[:k].reshape(-1).float().sum(0) * 100. / batch_size for k in topk]
33+
return [
34+
correct[:k].reshape(-1).float().sum(0) * 100. / batch_size
35+
if k <= maxk else torch.tensor(100.) for k in topk
36+
]

0 commit comments

Comments
 (0)