如何在PyTorch（或NumPy）中按特定稀疏性平铺二维张量模式？

我有一个二维Tensor模式，希望按特定稀疏性进行重复/平铺。我猜测PyTorch中存在某个函数（或使用gather、fold/unfold等的两行代码实现方案）可以完成此操作。

以下两种方法均可得到预期结果，但存在低效或操作晦涩的问题：

# inputs
mask = torch.rand(5, 5) > .7
pattern = torch.rand(4, 4)

# approach a (brute force):
result_a = torch.zeros(20, 20)
for i in range(5):
  for j in range(5):
    if mask[i][j]:
       result_a[i*4: (i+1)*4, j*4: (j+1)*4] = pattern

# approach b (interpolate the binary mask)
tiled_pattern = torch.tile(pattern, (5, 5))
bigger_mask = torch.nn.functional.interpolate(mask.view(1, 1, *mask.shape).to(torch.float16), scale_factor=(4, 4), mode='area').squeeze((0, 1))
result_b = torch.mul(tiled_pattern, bigger_mask)

我尤其不喜欢方法(b)，因为interpolate函数需要转换输入类型、添加并移除两个维度，而且我不确定'area'模式是否始终有效——PyTorch文档甚至没有明确定义该模式。

请问是否有更简洁的实现方式？若PyTorch中没有，NumPy（或其他常用Python包）是否可行？

PyTorch 实现 #

可以利用repeat_interleave直接将掩码扩展为目标尺寸，再与平铺后的模式相乘，逻辑直观且无需额外类型转换：

# 扩展掩码：每个元素在两个维度上分别重复pattern的高/宽次数
mask_expanded = mask.repeat_interleave(pattern.shape[0], dim=0).repeat_interleave(pattern.shape[1], dim=1)
# 平铺模式到目标尺寸
tiled_pattern = pattern.tile(mask.shape)
# 相乘得到结果
result = tiled_pattern * mask_expanded

原理说明：

• repeat_interleave(dim=0, repeats=4)会将原5x5掩码的每一行重复4次，得到20x5的中间结果；再对dim=1执行相同操作，最终得到20x20的掩码，每个原掩码元素对应4x4的区域，值与原元素一致。
• pattern.tile(mask.shape)等价于将4x4的pattern平铺5次行、5次列，得到20x20的模式矩阵。
• 两者相乘后，只有掩码为True的区域会保留模式值，其余区域为0，与原方法结果完全一致。

NumPy 实现 #

如果使用NumPy，可通过np.repeat实现类似逻辑：

import numpy as np

mask_np = mask.numpy()
pattern_np = pattern.numpy()

# 扩展掩码
mask_expanded_np = np.repeat(np.repeat(mask_np, pattern_np.shape[0], axis=0), pattern_np.shape[1], axis=1)
# 平铺模式
tiled_pattern_np = np.tile(pattern_np, mask_np.shape)
# 计算结果
result_np = tiled_pattern_np * mask_expanded_np

优势对比 #

• 相比方法a：避免了嵌套循环，计算效率更高（尤其是大尺寸张量场景）。
• 相比方法b：无需调整维度、转换数据类型，依赖的repeat_interleave/np.repeat逻辑明确，文档定义清晰，不存在行为不确定性。

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