pytorch中torch.masked_select()总结

• 1. 官方文档：

 

• 2. 当input 为 二维tensor时:

实例:

import torch

a = torch.randn(3,4)
mask = torch.ByteTensor([[1,1,1,0],[1,0,0,0],[0,0,0,0]])
b = torch.masked_select(a,mask)
print('a:',a)
print('mask:',mask)
print('b:',b) 

会发现,输入的tensor a 会根据 mask 对应位置的值 进行删选, 最后输出一个1维的tensor.

但此时出现了一个UserWarning: masked_select received a mask with dtype torch.uint8, this behavior is now deprecated,please use a mask with dtype torch.bool instead.

所以,可以进行如下改进:

import torch

a = torch.randn(3,4)
mask = torch.tensor([[1,1,1,0],[1,0,0,0],[0,0,0,0]],dtype=torch.bool)
b = torch.masked_select(a,mask)
print('a:',a)
print('mask:',mask)
print('b:',b) 

此时,就没有warning了.

 

注意: 官方文档说 mask 跟 input 有相同数量的元素数目, 但形状或维度不需要相同,     此处经实验发现有误, mask 和 input 需保证维度相同,否则会报错.

import torch

a = torch.randn(3,4)
mask = torch.tensor([[1,1,1,0,1,0],[0,0,0,0,0,0]],dtype=torch.bool)
b = torch.masked_select(a,mask)
print('a:',a)
print('mask:',mask)
print('b:',b) 

 

• 3. 当tensor 为多维tensor(比如三维)时:

比如, 考虑如下实际应用场景:  简单化, 二分类问题, 判断每个字后面的韵律是0还是1(实际上, 韵律要分#0, #1, #2, #3, #4这五种), 每个batch_size=3,即有3句话, 按照这三句话中的最长的句子进行padding成等长, 但是在计算loss时 要将padding的字进行mask掉, 得到真正的每个字

一般思路为: 先根据每个batch里的每句话的长度 得到一个二维的mask, 但是字符的tensor是三维的(其中第一维是句子数, 第二维是每句话的字数, 最后一维是每个字对应0或1的概率), 二者不匹配, 无法直接用函数 torch.masked_selected(input, mask) 进行mask, 所以, 要先用unsqueeze(-1)给mask增加第三维, 然后再用torch.masked_select()进行mask, 得到一维的概率值, 最后再用reshape(-1,2) 得到二维的tensor, 即每个字的属于0 或1的概率组成的tensor.

具体如下:

import torch

a = torch.randn(3,4,2)
print('a:',a)
print('a.shape:',a.shape)
mask1 = torch.tensor([[1,1,1,0],[1,0,0,0],[0,0,0,0]],dtype=torch.bool)
print('mask1:',mask1)
print('mask1.shape:',mask1.shape)
mask2 = mask1.unsqueeze(-1)
print('mask2:',mask2)
print('mask2.shape:',mask2.shape)
b1 = torch.masked_select(a,mask2)
print('b1:',b1)
print('b1.shape:',b1.shape)
b2 = b1.reshape(-1,2)
print('b2:',b2)
print('b2.shape:',b2.shape)