如何创建“岛屿”风格的pytorch矩阵

Question

可能有一个简单的问题，希望有一个简单的解决方案：

我得到了（稀疏）1D布尔张量[1，n]。

我想产生一个2D张量的大小[n，n]，其中包含由1D张量诱导的岛屿。观察以下图像示例是最容易的，其中鞋面是1D布尔量张量，下面的矩阵表示所得的矩阵：

Answer 1

给定一个蒙版输入：

>>> x = torch.tensor([0,0,0,1,0,0,0,0,1,0,0])

您可以用 ：

>>> index = x.nonzero()[:,0].diff(prepend=torch.zeros(1), append=torch.ones(1)*len(x))
tensor([3., 5., 3.])

然后使用代码> 创建对角线块矩阵：

>>> torch.block_diag(*[torch.ones(i,i) for i in index.int()])
tensor([[1., 1., 1., 0., 0., 0., 0., 0., 0., 0., 0.],
        [1., 1., 1., 0., 0., 0., 0., 0., 0., 0., 0.],
        [1., 1., 1., 0., 0., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 1., 1., 1., 1., 1., 0., 0., 0.],
        [0., 0., 0., 1., 1., 1., 1., 1., 0., 0., 0.],
        [0., 0., 0., 1., 1., 1., 1., 1., 0., 0., 0.],
        [0., 0., 0., 1., 1., 1., 1., 1., 0., 0., 0.],
        [0., 0., 0., 1., 1., 1., 1., 1., 0., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 1., 1., 1.],
        [0., 0., 0., 0., 0., 0., 0., 0., 1., 1., 1.],
        [0., 0., 0., 0., 0., 0., 0., 0., 1., 1., 1.]])