torch.LongTensor转换成one hot tensor编码

import torch
'''
pytorch0.4.0及以上的高版本，加上pytorch0.2.0都支持torch.Tensor._scatter
功能实现将torch.LongTensor编码成 one-hot vector的功能
'''
rois_label=torch.tensor(
    [[0,1,0,5,2,7,4,3,1,2]]
)#假设当前的ground truth boxes共有10个，类别范围从0-7，加上背景类别共有8个类别，背景类别用0表示
ids = rois_label.view(-1, 1)
cls_prob=torch.ones((10,8))
'''
cls_prob:表示网络模型所预测出来的当前region proposal（在RCNN中）
属于8个类别的概率值
'''
proposal_num=cls_prob.shape[0]
num_class=cls_prob.shape[1]
class_mask=cls_prob.data.new(proposal_num,num_class).fill_(0).cpu()
print('ids.data', type(ids.data), torch.min(ids), torch.max(ids))
class_mask.scatter_(1, ids.data.cpu(), 1.0)
print(class_mask.shape, class_mask)
'''
ids.data <class 'torch.Tensor'> tensor(0) tensor(7)
torch.Size([10, 8]) tensor([[1., 0., 0., 0., 0., 0., 0., 0.],
        [0., 1., 0., 0., 0., 0., 0., 0.],
        [1., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 1., 0., 0.],
        [0., 0., 1., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 1.],
        [0., 0., 0., 0., 1., 0., 0., 0.],
        [0., 0., 0., 1., 0., 0., 0., 0.],
        [0., 1., 0., 0., 0., 0., 0., 0.],
        [0., 0., 1., 0., 0., 0., 0., 0.]])
'''

'''
pytorch 0.3.0 版本的 torch.Tensor._scatter函数并不能使用
故而需要用户自己编写函数实现对于输入进行one-hot编码的功能
'''
class_mask2=cls_prob.data.new(proposal_num,num_class).fill_(0).cpu()
matrix=torch.eye(num_class, device='cpu')  # [D,D]
class_mask2=matrix[rois_label]
print(class_mask2.type(),class_mask.type())
if torch.sum(class_mask2-class_mask)==0:
    print('true')
'''
torch.FloatTensor torch.FloatTensor
true
'''

注意matrix[rois_label]这种索引方式的索引值必须是torch.LongTensor，且必须放置在CPU上，不能是GPU。我运行程序时将rois_label转换成了variable类型，故而会报错。

debug方法

class_mask=matrix[rois_label.data.cpu()]

首先通过  .data   将Variable转换成torch.Tensor类型，然后再将cuda上的tensor放置到CPU上面。