如何从形状(1,1,128,128,128)的标签张量获取索引以转One-Hot？

Hey there! Let's clear up this confusion for you—your existing label tensor is actually already what you need as an index tensor, we just need to make a tiny adjustment to fit torch.nn.functional.one_hot's requirements.

Here's the breakdown and step-by-step solution:

Key Notes First #

• torch.nn.functional.one_hot expects integer-type tensors where each value represents an index starting from 0.
• Your label tensor has values from 1 to 24, so we just need to shift these values down by 1 to get 0-based indices (0 to 23), which perfectly maps to 24 classes.

Step-by-Step Code #

import torch

# Let's assume your original label tensor looks like this (shape: (1,1,128,128,128), values 1-24)
label_tensor = torch.randint(low=1, high=25, size=(1,1,128,128,128))

# 1. Ensure the tensor is of integer type (required by one_hot)
label_tensor = label_tensor.long()

# 2. Convert to 0-based index tensor by subtracting 1
index_tensor = label_tensor - 1

# 3. Apply one-hot encoding with 24 classes
one_hot = torch.nn.functional.one_hot(index_tensor, num_classes=24)

# Optional: Rearrange dimensions if needed (e.g., move class channel to the front)
# Original one_hot shape: (1,1,128,128,128,24)
# After permute: (1,24,1,128,128,128) – which is more common for image segmentation tasks
one_hot = one_hot.permute(0, 5, 1, 2, 3, 4)

Why This Works #

• By subtracting 1, we turn your label values (1→24) into valid 0-based indices (0→23). This matches how one_hot expects indices: each index corresponds to a position in the one-hot vector.
• The long() conversion ensures the tensor is of the correct integer type (PyTorch's one-hot function doesn't work with float tensors).

If you actually intended to include 0 as a valid class (you mentioned a range of 0-24 initially), just adjust num_classes to 25 and skip the subtraction step—your original tensor would already be a valid index tensor in that case.

内容的提问来源于stack exchange，提问作者Ryan