32 lines
801 B
Python
32 lines
801 B
Python
import torch.nn as nn
|
|
|
|
|
|
def accuracy(pred, target, topk=1):
|
|
assert isinstance(topk, (int, tuple))
|
|
if isinstance(topk, int):
|
|
topk = (topk, )
|
|
return_single = True
|
|
else:
|
|
return_single = False
|
|
|
|
maxk = max(topk)
|
|
_, pred_label = pred.topk(maxk, dim=1)
|
|
pred_label = pred_label.t()
|
|
correct = pred_label.eq(target.view(1, -1).expand_as(pred_label))
|
|
|
|
res = []
|
|
for k in topk:
|
|
correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
|
|
res.append(correct_k.mul_(100.0 / pred.size(0)))
|
|
return res[0] if return_single else res
|
|
|
|
|
|
class Accuracy(nn.Module):
|
|
|
|
def __init__(self, topk=(1, )):
|
|
super().__init__()
|
|
self.topk = topk
|
|
|
|
def forward(self, pred, target):
|
|
return accuracy(pred, target, self.topk)
|