r"""Calculates the classification accuracy given predicted logits and ground-truth labels. Args: logits: model predictions of shape `[batch_size, num_classes]`, representing the probability (likelyhood) of each class. target: ground-truth labels, 1d tensor of int32.
(
logits: Tensor, target: Tensor, topk: Union[int, Iterable[int]] = 1
)
| 14 | |
| 15 | |
| 16 | def topk_accuracy( |
| 17 | logits: Tensor, target: Tensor, topk: Union[int, Iterable[int]] = 1 |
| 18 | ) -> Union[Tensor, Iterable[Tensor]]: |
| 19 | r"""Calculates the classification accuracy given predicted logits and ground-truth labels. |
| 20 | |
| 21 | Args: |
| 22 | logits: model predictions of shape `[batch_size, num_classes]`, |
| 23 | representing the probability (likelyhood) of each class. |
| 24 | target: ground-truth labels, 1d tensor of int32. |
| 25 | topk: specifies the topk values, could be an int or tuple of ints. Default: 1 |
| 26 | |
| 27 | Returns: |
| 28 | tensor(s) of classification accuracy between 0.0 and 1.0. |
| 29 | """ |
| 30 | if isinstance(topk, int): |
| 31 | topk = (topk,) |
| 32 | _, pred = _topk(logits, k=max(topk), descending=True) |
| 33 | accs = [] |
| 34 | for k in topk: |
| 35 | correct = pred[:, :k].detach() == broadcast_to( |
| 36 | transpose(target, (0, "x")), (target.shape[0], k) |
| 37 | ) |
| 38 | accs.append(correct.astype(np.float32).sum() / target.shape[0]) |
| 39 | if len(topk) == 1: # type: ignore[arg-type] |
| 40 | accs = accs[0] |
| 41 | return accs |