pythonnumpypytorch

Tile according to a pattern in pytorch (or numpy)?


I have a 2-d tensor pattern I'd like to repeat/tile with a particular sparsity. I suspect there is a function (or two-line approach using gather, fold/unfold, or similar) to do this in Pytorch.

These two approaches each get the result I want but are inefficient and/or confusing:

# 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)

In particular, I don't like approach (b) because the interpolate function needs me to cast the input, add and remove two dimensions, and I don't trust the 'area' mode to always work--it's not even really defined in the Pytorch documentation.

Is there a cleaner way to do this? If not in Pytorch, than in Numpy (or another common Python package?)


Solution

  • You can do this simply with a broadcasted multiply, in both torch and numpy.

    result_c = (mask[:, None, :, None] * pattern[None, :, None, :]).reshape(20, 20)