1 torch.Tensor.masked_fill参数详解与使用

1.1 torch.Tensor.masked_fill参数详解

1. 函数形式

torch.Tensor.masked_fill(mask, value)

2. 函数功能
输入的mask需要与当前的基础Tensor的形状一致。
将mask中为True的元素对应的基础Tensor的元素设置为值value。

3. 函数参数

  • mask：mask既可以是int型Tensor（值为0或者1）也可以是bool型Tensor（值为False或者True）
  • value：float，填充的值

4. 函数返回值
返回填充后的Tensor

1.2 torch.Tensor.masked_fill使用示例

下面一个简单的例子说明masked_fill函数的使用，首先我们创建一个4x4的一个基础矩阵，然后创建一个4x4的对角矩阵，然后根据对角矩阵将对角线上的基础机矩阵的值全部设置为100，具体的代码如下所示。

import torch

if __name__ == '__main__':
    tensor = torch.arange(0,16).view(4,4)
    print('origin tensor:\n{}\n'.format(tensor))

    mask = torch.eye(4,dtype=torch.bool)
    print('mask tensor:\n{}\n'.format(mask))

    tensor = tensor.masked_fill(mask,100)
    print('filled tensor:\n{}'.format(tensor))

输出

origin tensor:
tensor([[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11],
        [12, 13, 14, 15]])

mask tensor:
tensor([[ True, False, False, False],
        [False,  True, False, False],
        [False, False,  True, False],
        [False, False, False,  True]])

filled tensor:
tensor([[100,   1,   2,   3],
        [  4, 100,   6,   7],
        [  8,   9, 100,  11],
        [ 12,  13,  14, 100]])