【最佳实践】pytorch获取top1和topk准确率

def eval(eval_dataloader,k = 5):
    with torch.no_grad():
        total = 0
        top1 = 0
        topk = 0
        for (test_imgs, test_labels) in eval_dataloader:
            test_labels = test_labels.to(device)

            preds = pred_net(linear_classify_net(backbone_net(test_imgs.to(device))))

            _,maxk = torch.topk(preds,k,dim=-1)
            total += test_labels.size(0)
            test_labels = test_labels.view(-1,1) # reshape labels from [n] to [n,1] to compare [n,k]

            top1 += (test_labels == maxk[:,0:1]).sum().item()
            topk += (test_labels == maxk).sum().item()

        print('Accuracy of the network on total {} test images: @top1={}%; @top{}={}%'.format(total,100 * top1 / total,k,100*topk/total))