Batched sub-matrix selection

1

Is it possible to select a batch of sub-matrices with different starting and ending rows and column indices from a matrix? Please check the below example:

# construct a matrix
m = torch.arange(0, 9).view(3, 3)
# define starting and ending indices for each batch
start_row = torch.tensor([[0,0], [1, 0]])
end_row = torch.tensor([[1, 1], [2, 1]])
start_col = torch.tensor([[0,0], [0,0]])
end_col = torch.tensor([[1,1], [1,1]])
# selecting sub-matrices
print(m[start_row:end_row, start_col:end_col])
2

Looks like advanced indexing. I bet converting this to numpy and asking on stackoverflow will get answers fast, numpy questions usually get answered pretty fast.

3

Have you find an answer to your question? I look for this solution as well

4

PyTorch doesn’t have a way to vectorize this computation. You could use masking, but that’s going to likely be quite inefficient as well.

For readers who come here later, the naive version of what the question is asking for is this:

def naive_batch_submatrix(m, start_row, end_row, start_col, end_col):
    samples = []
    for i in range(start_row.shape[0]): # number of batches
        # for each batch
        batch = []
        for j in range(start_row.shape[1]):        
            batch.append(m[start_row[i, j]:end_row[i, j]+1, start_col[i, j]:end_col[i, j]+1])
        samples.append(batch)
    return samples

One option is to torch.compile this naive version, and see if that helps. For these small matrix sizes and a small number of start_row / end_row it might not matter, but for larger-scale it likely will be more efficient and remove a bunch of overhead.

5

Hi, actually I find a way to select sub matrices in vectorize way.
asumme you have a start_x vector of different indexes in size (batch,) and you have a start_y vector of different indexes in size (batch,).
now I use the idea of numpy advance indexing from here: numpy doc

first create a map of all the indexing I interested in with size (batch_size, step_size, step_size)
(in my example I use a window around the center that I detect with argmax)

start_x = target_mat_inx[:, 0] - nn_cfg.window_size_around_center # shape(batch_size,)
start_y = target_mat_inx[:, 1] - nn_cfg.window_size_around_center # shape(batch_size,)
start_x_mat = (start_x.reshape(-1, 1) + torch.arange(2*nn_cfg.window_size_around_center, device=device)).unsqueeze(-1).repeat(1, 1, 2*nn_cfg.window_size_around_center) # shape(batch_size, 2* step_size, 2* step_size)
start_y_mat = (start_y.reshape(-1, 1) + torch.arange(2*nn_cfg.window_size_around_center, device=device)).unsqueeze(1).repeat(1, 2*nn_cfg.window_size_around_center, 1)  # shape(batch_size, 2* step_size, 2* step_size)

now, create the batch selection vector:

batch_index = torch.arange(nn_cfg.batch_size, device=device).reshape(-1,1,1)

select the sub matrixes in all batch

matrix_mask_target = matrix[batch_index, start_x_mat, start_y_mat] # shape(batch_size, 2* step_size, 2* step_size)

note: it is possible to select sub bathes as well, just change the batch_index to specific selection and make sure that start_x and start_y align with your selection