pytorch的tensor张量如何逐像素比较计算

Question

我有一个网络输出的tensor张量矩阵z，现在想对这个矩阵做如下约束：
矩阵大小为BCH*W，理想状况下，每行每个元素H（i）应该大于之前的元素（H0~H（i-1）），设为1；如果不是最大，则设为0.
我将这个算法转到numpy实现了一下，但不知道如何直接在tenosr下实现，因为在numpy下无法反向传播回去。

def cal_shadow(h,z,x_ij):
    tan_sigma=x_ij/(h-z)
    tan_sigma=tan_sigma.cpu().detach().numpy()
    shadow=np.copy(tan_sigma)
    for k in range(tan_sigma.shape[0]):
        for i in range(tan_sigma.shape[2]):
            max_sigma = -100000
            for j in range(tan_sigma.shape[3]):
                if tan_sigma[k, 0, i, j] > max_sigma:
                    shadow[k, 0, i, j] = 1
                    max_sigma = tan_sigma[k, 0, i, j]
                else:
                    shadow[k, 0, i, j] = 0 #阴影都设为0
    shadow = torch.from_numpy(shadow)
    shadow = shadow.type(torch.FloatTensor)
    shadow=shadow.to(z.device)
    return shadow

Answer 1

import torch

def cal_shadow(h, z):
    # 计算每个元素和之前的元素的最大值
    cummax = torch.cummax(z, dim=-1).values
    # 创建一个大小相同的零张量
    shadow = torch.zeros_like(z)
    # 如果元素大于等于之前的最大值，则将其设置为1，否则为0
    shadow[z >= cummax] = 1
    # 将张量类型转换为与输入张量相同的类型，并将其发送到相同的设备
    shadow = shadow.type(z.dtype).to(z.device)
    return shadow