PyTorch中如何根据掩码值合并两张图像（附NumPy实现示例）

核心合并规则和给出的NumPy逻辑完全一致：掩码值为1（布尔True）的位置取第一张图像像素，其余位置取第二张图像像素。PyTorch原生支持布尔掩码索引，同时提供了更适合张量运算的向量化实现。

基础实现（和NumPy写法一一对应） #

PyTorch中用clone()方法替代NumPy的copy()做张量深拷贝，布尔掩码的索引赋值逻辑和NumPy完全一致，适用于和示例相同的(H, W, C)（高、宽、通道）维度顺序的图像张量：

import torch

# 生成测试数据，维度顺序和NumPy示例一致
img1 = torch.rand(100, 100, 3)
img2 = torch.rand(100, 100, 3)
mask = torch.rand(100, 100) > 0.5  # 布尔掩码，形状为(H,W)

res = img2.clone()
res[mask] = img1[mask]

适配PyTorch常规图像格式的写法 #

PyTorch视觉任务中通常使用(C, H, W)（通道在前）的维度顺序存储图像，直接用上面的写法会出现维度不匹配问题，只需要在索引时指定通道维度做广播即可：

# 通道在前格式的测试数据
img1 = torch.rand(3, 100, 100)
img2 = torch.rand(3, 100, 100)
mask = torch.rand(100, 100) > 0.5

res = img2.clone()
# 冒号表示选中所有通道，掩码自动广播到3个通道
res[:, mask] = img1[:, mask]

训练场景推荐：无额外拷贝的向量化实现 #

如果是在模型训练流程中使用，推荐用torch.where实现，不需要手动做clone拷贝，对自动求导更友好，运行效率更高：

# 给掩码增加1个通道维度，从(H,W)变为(1,H,W)，自动匹配图像的(C,H,W)形状
mask_broadcast = mask.unsqueeze(0)
# 规则：mask为True的位置取img1，否则取img2
res = torch.where(mask_broadcast, img1, img2)

注意事项 #

• 所有参与运算的img1、img2、mask必须在同一个设备上（同CPU或同GPU），否则会触发设备不匹配报错
• 掩码必须是布尔类型（torch.bool），如果掩码是0/1数值型，用mask = mask.to(torch.bool)转换即可
• 如果处理批量图像（张量形状为(B, C, H, W)，B为批次大小），只需要把掩码调整为(B, 1, H, W)的形状，上面的torch.where逻辑可以直接批量运行，不需要写循环

内容的提问来源于stack exchange，提问作者John M.