操作数组问题

        a=torch.rand(2,2,3,4)

        a[a>0.5] = 1  ###a>0.5
        a[a<=0.5] = 0
        print(a)
        index = torch.nonzero(a==0)
        index = index[index[:,3] != (a.size(3)-1)]
        index[:,3] = index[:,3]+1
        a[index[:,0], index[:,1], index[:,2], index[:,3]] = 0
        print(a)




        a=torch.tensor(range(24)).view(2,2,3,2)
        b=torch.tensor(range(24)).view(2,2,3,2)
        c = torch.rand(2,2,3,2)
        c[c>0.5] = 1  ###a>0.5
        c[c<=0.5] = 0
        b[c.byte()]

这里的b和c是一样的维度,所以b[c.byte()]就是把矩阵c中为1位置上的b的值取出来,b[c.byte()]的维度为[10],表示有10个满足条件的值

猜你喜欢

转载自blog.csdn.net/Gussss/article/details/102245040