如何在PyTorch张量中按多区间条件赋值？

针对你的需求，PyTorch不支持直接写0 <= x <=1这类连续区间判断，但可以通过链式调用torch.where或布尔掩码赋值两种方式实现，以下是具体实现：

方法1：链式使用torch.where #

利用torch.where的执行顺序，按条件依次处理（后一步会覆盖前一步未匹配的元素）：

import torch

# 初始化你的张量
x = torch.tensor([[ 0.2213, -0.1180,  1.1186],
                  [-0.9943, -0.7679, -1.7057]])

# 按条件优先级处理：先处理负数，再处理0-1区间，最后处理2-3区间
output = torch.where(x < 0, torch.tensor(0.0), x)
output = torch.where((0 <= output) & (output <= 1), torch.tensor(5.0), output)
output = torch.where((2 <= output) & (output <= 3), torch.tensor(7.0), output)

print(output)

执行后输出：

tensor([[5.0000, 0.0000, 1.1186],
        [0.0000, 0.0000, 0.0000]])

方法2：布尔掩码直接赋值 #

通过创建元素级的布尔掩码，针对性地为对应区间赋值：

import torch

x = torch.tensor([[ 0.2213, -0.1180,  1.1186],
                  [-0.9943, -0.7679, -1.7057]])

# 创建各区间的掩码
mask_negative = x < 0
mask_0_to_1 = (0 <= x) & (x <= 1)
mask_2_to_3 = (2 <= x) & (x <= 3)

# 克隆原张量作为输出容器，再按掩码赋值
output = x.clone()
output[mask_negative] = 0
output[mask_0_to_1] = 5
output[mask_2_to_3] = 7

print(output)

关键说明 #

• 注意PyTorch中必须拆分连续区间判断：把0 <= x <=1拆成(0 <= x) & (x <=1)，用按位与&连接两个布尔张量，同时括号不能省略（避免运算符优先级问题）。
• 如果存在区间重叠的情况，需要注意赋值顺序：后执行的赋值会覆盖之前的结果，所以要按业务逻辑确定优先级。

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