torchmetrics.utilities.data

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)

Convert 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])

Convert 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]])