Get a 10x10 patch from a 100x100 pytorch tensor with torus style wrap around the boundries

user_2341735 0
Viewed 119

How can I get a 10x10 patch from a 100x100 pytorch tensor, with the added constraint that if a patch would go outside the boundaries of the array, then it wraps around the edges (as if the array was a torus, with the top joined to the bottom, and the left joined to the right)?

I wrote this code that will do the job, I'm looking for something more elegant, efficient and clear:

def shift_matrix(a, distances) -> Tensor:
  x, y = distances
  a = torch.cat((a[x:], a[0:x]), dim=0)
  a = torch.cat((a[:, y:], a[:, :y]), dim=1)
  return a

def randomly_shift_matrix(a) -> Tensor:
  return shift_matrix(a, np.random.randint(low = 0, high = a.size()))

def random_patch(a, size) -> Tensor:
  full_shifted_matrix = randomly_shift_matrix(a)
  return full_shifted_matrix[0:size[0], 0:size[1]]

I feel like something with negative index slices should work. I haven't found it though.

You can see the code in google colab here.

pytorch slice arrays python
1 Answers

You are looking for torch.roll

def random_patch(a, size) -> Tensor:
  shifts = np.random.randint(low = 0, high = a.size())
  return torch.roll(a, shifts=shifts, dims=(0, 1))[:size[0], :size[1]]
user_17144100
Related