PyTorch下如何基于分割张量C将张量A对应区域复制到全白张量B

核心逻辑直接用布尔掩码广播赋值即可，无需使用torch.masked_select，该方法会保留原张量的维度结构，完全匹配你的需求。

实现步骤 #

• 第一步：指定你需要提取的目标分割类别ID，构造对应位置的布尔掩码
• 第二步：将掩码扩展匹配图像张量的维度，适配batch和3通道维度
• 第三步：直接对全白张量B的掩码位置赋值为原图张量A的对应位置值

完整代码示例 #

import torch

# 按需修改为你要提取的分割类别ID，示例为提取类别37的区域
target_class = 37

# 1. 生成布尔掩码：对应目标类别位置为True，其余为False
mask = (C == target_class)
# 若需要同时提取多个类别，替换为下行代码即可：
# mask = torch.isin(C, torch.tensor([337, 340], device=C.device))

# 2. 扩展掩码维度，适配A/B的(1,3,256,256)维度，支持广播
mask = mask.unsqueeze(0).unsqueeze(0)

# 3. 赋值操作：将A对应区域的内容复制到B的对应位置
B[mask.expand_as(B)] = A[mask.expand_as(B)]

注意事项 #

请确保张量A、B、C在同一设备上（如都在cuda:0或者都在cpu），否则会触发设备不匹配报错。
该逻辑可通用适配不同batch、不同尺寸的输入，仅需根据实际维度调整掩码的扩展规则即可。

内容的提问来源于stack exchange，提问作者AlwaysNull