如何在3D张量中忽略填充行计算均值

我有一个形状为(batch_size, N, dim)的3D张量A，以及一个形状为(batch_size, N, 2)的3D张量B。B中存在填充行以补至长度N（这些填充行并非零向量，因为已通过一些函数处理）。判断某行是否为填充行需查看张量B：若第k行是填充行，则B中第k行的值为[0, 0]。我希望在计算均值前过滤掉A中的这些填充行，最终得到形状为(batch_size, dim)的均值结果。

补充：我已想到一种解决方案，欢迎提供其他方案！

我的解决方案 #

# 压缩为2D张量
A = A.view(-1, A.size(-1))
B = B.view(-1, B.size(-1))

# 生成填充行掩码
mask = torch.sum(B, dim=-1)
mask[mask!=0] = 1

# 将A中的填充行置零
A = mask.unsqueeze(-1)*A

# 计算每个batch的有效样本数
mask = mask.view(batch_size, -1)
batch_cluster_size = torch.sum(mask, dim=-1, keepdim=True)

# 恢复原形状
A = A.view(batch_size, -1, A.size(-1))

output = torch.sum(A, dim=1)/batch_cluster_size

其他解决方案 #

方案一：3D张量直接操作（更严谨直观） #

# 生成掩码：B的行全为[0,0]时为False，否则为True
mask = ~torch.all(B == 0, dim=-1)  # 形状: (batch_size, N)

# 计算每个batch的有效样本数，避免除以0（全填充时设为1）
batch_cluster_size = mask.sum(dim=1, keepdim=True).clamp(min=1)  # 形状: (batch_size, 1)

# 掩码扩展到与A同维度，过滤后求和取均值
mask_expanded = mask.unsqueeze(-1)  # 形状: (batch_size, N, 1)
filtered_sum = (A * mask_expanded).sum(dim=1)  # 形状: (batch_size, dim)
output = filtered_sum / batch_cluster_size

这个方案无需展平张量，直接在原维度处理，逻辑更清晰。同时用torch.all(B == 0)判断填充行更严谨——原方案通过求和判断可能遇到非零元素和为0的极端情况（比如B行是[1,-1]），而题目明确填充行是[0,0]，用all判断完全匹配需求。

方案二：用torch.where替换填充行 #

# 标记填充行：B的行全为[0,0]时为True
mask = torch.all(B == 0, dim=-1)  # 形状: (batch_size, N)

# 将填充行对应的A元素替换为0
A_filtered = torch.where(mask.unsqueeze(-1), torch.zeros_like(A), A)

# 计算有效样本数，避免除以0
batch_cluster_size = (~mask).sum(dim=1, keepdim=True).clamp(min=1)

# 计算均值
output = A_filtered.sum(dim=1) / batch_cluster_size

该方案通过torch.where直接完成填充行的替换，代码可读性高，同样不需要张量形状转换的额外操作。

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