Python中不使用循环对numpy/torch张量数组批量应用掩码的方法

核心原理是利用广播机制对齐掩码和输入数组的维度，无需显式循环，底层自动完成向量化运算，效率远高于Python级别的for循环。

PyTorch 实现 #

你的输入张量arr形状为(N, M, H, W)，掩码mask形状为(N, H, W)，只需要给掩码新增一个对应M维度的轴即可完成广播逐元素相乘：

# 方法1：用unsqueeze新增第1维（维度索引从0开始）
result = arr * mask.unsqueeze(1)

# 方法2：用None索引新增维度，写法更简洁
result = arr * mask[:, None, :, :]

运算时形状为(N, 1, H, W)的掩码会自动广播到和arr一致的(N, M, H, W)，每一组M个矩阵都会复用对应索引的掩码，和你写的循环逻辑完全等价。

Numpy 实现 #

逻辑和PyTorch完全一致，只需对齐维度即可：

# 方法1：用expand_dims新增轴
result = arr * np.expand_dims(mask, axis=1)

# 方法2：用np.newaxis新增维度
result = arr * mask[:, np.newaxis, ...]

常见问题说明 #

你之前调用np.dot、np.matmul报错是因为这两个API是做矩阵乘法运算，而你的需求是逐元素相乘，不需要做矩阵乘逻辑；np.multiply报错是因为没有对齐两个输入的维度，只要按照上面的方法给掩码加一个维度即可正常调用：

# numpy下用multiply的正确写法
result = np.multiply(arr, np.expand_dims(mask, 1))

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