Batched sub-matrix selection

Junhong: <p>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:</p>
<pre><code class="lang-auto"># 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])
</code></pre>

Santosh_Gupta: <p>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.</p>

AviranRSX: <p>Have you find an answer to your question? I look for this solution as well</p>

smth: <p>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.</p>
<p>For readers who come here later, the naive version of what the question is asking for is this:</p>
<pre><code class="lang-auto">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
</code></pre>
<p>One option is to <code>torch.compile</code> this naive version, and see if that helps. For these small matrix sizes and a small number of <code>start_row</code> / <code>end_row</code> it might not matter, but for larger-scale it likely will be more efficient and remove a bunch of overhead.</p>

AviranRSX: <p>Hi, actually I find a way to select sub matrices in vectorize way.<br>
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,).<br>
now I use the idea of numpy advance indexing from here: <a href="https://numpy.org/doc/stable/user/basics.indexing.html#advanced-indexing" rel="noopener nofollow ugc">numpy doc</a></p>
<p>first create a map of all the indexing I interested in with size (batch_size, step_size, step_size)<br>
(in my example I use a window around the center that I detect with argmax)</p>
<pre><code class="lang-auto">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)
</code></pre>
<p>now, create the batch selection vector:</p>
<pre><code class="lang-auto">batch_index = torch.arange(nn_cfg.batch_size, device=device).reshape(-1,1,1)
</code></pre>
<p>select the sub matrixes in all batch</p>
<pre><code class="lang-auto">matrix_mask_target = matrix[batch_index, start_x_mat, start_y_mat] # shape(batch_size, 2* step_size, 2* step_size)
</code></pre>
<p>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</p>