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

我有一个网络输出的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

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

该回答通过自己思路及引用到GPTᴼᴾᴱᴺᴬᴵ搜索,得到内容具体如下。

您可以尝试使用PyTorch的torch.where和torch.cummax函数实现这个约束。以下是一个示例实现：

import torch

def cal_shadow(h, z, x_ij):
    tan_sigma = x_ij / (h - z)
    
    # 计算每个位置的累计最大值
    cum_max, _ = torch.cummax(tan_sigma, dim=-1)
    
    # 将tan_sigma与cum_max逐元素比较，如果tan_sigma大于cum_max中的前一个元素，则设为1，否则设为0
    shadow = torch.where(tan_sigma > torch.cat([torch.zeros_like(cum_max[..., :1]), cum_max[..., :-1]], dim=-1), 1, 0)
    
    return shadow.type(torch.FloatTensor)

这个实现中，我们首先计算了tan_sigma，然后使用torch.cummax函数计算每个位置的累计最大值。接下来，我们使用torch.where逐元素比较tan_sigma与cum_max中的前一个元素，如果tan_sigma大于前一个元素，则设为1，否则设为0。

这个实现仅使用了PyTorch的张量操作，因此可以在计算图中保留梯度信息，从而实现反向传播。

如果以上回答对您有所帮助，点击一下采纳该答案～谢谢

• 这篇博客也许可以解决你的问题👉 ：pytorch将tensor转为numpy用于图像绘制
• 除此之外, 这篇博客: pytorch中tensor的加减乘除和常见操作中的 从numpy中获得数据 部分也许能够解决你的问题, 你可以仔细阅读以下内容或跳转源博客中阅读:
  

• torch.from_numpy(ndarry)
  

  注：生成返回的tensor会和ndarry共享数据，任何对tensor的操作都会影响到ndarry,
  反之亦然

   a = numpy.array([1, 2, 3])
  
   t = torch.from_numpy(a)
  
  t
  
  tensor([ 1,  2,  3])
  
  t[0] = -1
  
   a
  
  array([-1,  2,  3])