torchmetrics.utilities.data

select_topk

torchmetrics.utilities.data.select_topk(prob_tensor, topk=1, dim=1)

Convert a probability tensor to binary by selecting top-k the highest entries.

Parameters

• prob_tensor (Tensor) – dense tensor of shape [..., C, ...], where C is in the position defined by the dim argument

• topk (int) – number of the highest entries to turn into 1s

• dim (int) – dimension on which to compare entries

Return type

Tensor

Returns

A binary tensor of the same shape as the input tensor of type torch.int32

Example

>>> x = torch.tensor([[1.1, 2.0, 3.0], [2.0, 1.0, 0.5]])
>>> select_topk(x, topk=2)
tensor([[0, 1, 1],
        [1, 1, 0]], dtype=torch.int32)

to_categorical

torchmetrics.utilities.data.to_categorical(x, argmax_dim=1)

Converts a tensor of probabilities to a dense label tensor.

Parameters

• x (Tensor) – probabilities to get the categorical label [N, d1, d2, …]

• argmax_dim (int) – dimension to apply

Return type

Tensor

Returns

A tensor with categorical labels [N, d2, …]

Example

>>> x = torch.tensor([[0.2, 0.5], [0.9, 0.1]])
>>> to_categorical(x)
tensor([1, 0])

to_onehot

torchmetrics.utilities.data.to_onehot(label_tensor, num_classes=None)

Converts a dense label tensor to one-hot format.

Parameters

• label_tensor (Tensor) – dense label tensor, with shape [N, d1, d2, …]

• num_classes (Optional[int]) – number of classes C

Return type

Tensor

Returns

A sparse label tensor with shape [N, C, d1, d2, …]

Example

>>> x = torch.tensor([1, 2, 3])
>>> to_onehot(x)
tensor([[0, 1, 0, 0],
        [0, 0, 1, 0],
        [0, 0, 0, 1]])