使用NumPy或OpenCV应用掩码是一个相对简单的过程。然而，如果我需要在优化算法的损失计算中使用掩码图像，我需要专门使用PyTorch，因为其他方式会干扰梯度计算。假设我有一个图像张量[1, 512, 512, 3]（批次、高度、宽度、通道）和一个掩码张量[1, 20, 512, 512]（批次、通道、高度、宽度），其中每个通道对应于20个分割类别中的一个，我希望得到一个掩码图像张量，除了属于一个或多个指定分割类别的像素外，其他像素都填充为黑色(0, 0, 0)。

以下是使用numpy完成此操作的方法：

import numpy as npimport torch# 创建虚拟图像和掩码image_tensor = torch.randn([1, 512, 512, 3])mask_tensor = torch.randn([1, 20, 512, 512])# 对掩码应用argmaxmask_tensor =  torch.max(mask_tensor, 1)[1] # -> 1, 512, 512# 定义掩码函数def selective_mask(image_src, mask, dims=[]):    h, w = mask.shape    background = np.zeros([h, w, 3], dtype=np.uint8)    for j_, j in enumerate(mask[:, :]):        for k_, k in enumerate(j):            if k in dims:                background[j_, k_] = image_src[j_, k_]    output = background    return output# 将张量转换为numpy:image = image_tensor.squeeze(0).cpu().numpy()mask = mask_tensor.squeeze(0).cpu().nmpy()# 对多个类别应用掩码函数image_masked = selective_mask(image, mask, dims=[5, 6, 8])

我的代码应该如何修改以符合PyTorch的要求？

回答：

首先，selective_mask函数的定义远非您所说的“简单直接”。使用numpy（以及设计上与之大多兼容的torch）的关键点是利用操作的向量化，并避免使用循环，因为循环无法并行化。

如果您以这种方式重写上述函数：

def selective_mask(image_src, mask, channels=[]):    mask = mask[np.array(channels).astype(int)]    return np.sign(np.sum(mask, axis=0), dtype=image_src.dtype) * image_src

您会发现实际上可以用PyTorch张量做同样的事情（这里不需要压缩批次（第一个）维度）：

def selective_mask_t(image_src, mask, channels=[]):    mask = mask[:, torch.tensor(channels).long()]    mask = torch.sgn(torch.sum(mask, dim=1)).to(dtype=image_src.dtype).unsqueeze(-1)    return mask * image_src

此外，您可能希望以这种方式生成掩码本身：
（顺便说一句，这里结合使用max和sgn实际上比设置argmax索引的元素更快）

# 创建虚拟图像和掩码image_tensor = torch.randn([1, 512, 512, 3])mask_tensor = torch.randn([1, 20, 512, 512])# 离散化掩码（在具有最高值的通道中设置为1） -> 1, 20, 512, 512mask_tensor = torch.sgn(mask_tensor - torch.max(mask_tensor, 1)[0].unsqueeze(1)) + 1.  

然后它应该可以正常工作：

print(selective_mask_t(image_tensor, mask_tensor, [5, 6, 8]))