One-hot的使用

Posted 2023-04-01 MartinRY

pytorch官方已经提供了具体one-hot函数，可以直接使用

from torch.nn.functional import one_hot

result = one_hot(data,num_class)

首先介绍下，one_hot可以运用在当调用损失函数时，所放入数据尺寸不等的时候

例如：

    raise ValueError("Target size () must be the same as input size    ()".format(target.size(), input.size()))

ValueError: Target size (torch.Size([2, 256, 256])) must be the same as input size (torch.Size([2, 2, 256, 256]))

可以看到上面输入的预测是一个[2, 2, 256, 256] 的tensor（分别是nchw，即batch_size, channels, height, width）

而label的则是[2, 256, 256] ，缺少了通道数(因为label是单通道)

因此可以使用one_hot进行转换,将tensor数据直接放入one_hot中的第一个参数的位置,因为只有2类，所以num_class=2

from torch.nn.functional import one_hot

result = one_hot(label,2)

但直接放有时会遇到数据类型的问题，由于one_hot只接受int类型，所以若label数据是float类型，需要先将类型转换为int

label = label.to(dtype=torch.int64)

然后再放进one_hot

最后，由于模型的bias比较常为float类型，所以直接把输出的one_hot放进去会报数据类型错误

，因此在one_hot输出后，还需要再转换为float类型

one_hot_output = one_hot_output.type(torch.float32)

最后，把one_hot输出和pred放进计算loss的函数进行计算