Efficient and Fast way to traverse an array and compute mean

username: Yasser_Khalil
cooked: <p>Hello,</p>
<p>I have a problem that I solved. However, my solution is extremely slow. If someone could help me speed up the solution I would be extremely grateful.</p>
<p>So I am trying to compute the average of certain elements in an array using indices of another array. Please note that the same indices may exist multiple times. The average is saved in a smaller dimension array.<br>
Here is an illustrative example:</p>
<pre><code class="lang-auto">result = torch.full(( 32, 256, 256), -1).to(device)
data = torch.full(( 32, 40000), 1).to(device) #example of 
count = np.zeros(256,256)
indices = np.array([[1,2],[0,0],[1,2],[1,1],[1,0],[0,0],[1,2],[2,2],[2,2],[0,0]])

for i, indx in enumerate(indices):
        if(count[indx [0], indx [1]] == 0):
            result[:, indx [0], indx [1]] = data [:, i]
        else:
            result[:, indx [0], indx [1]] += data [:, i]
        count[indx [0], indx [1]] += 1

mask_zero = (count == 0)
count[mask_zero] = 1
count = torch.Tensor(count).to(device)
average= result/count

</code></pre>
<p>Is it possible to speed this?<br>
This is done for only one batch, is it even possible to generalize it to multiple batches, where the sizes of the same arrays are like:</p>
<pre><code class="lang-auto">result = torch.full((batch, 32, 256, 256), -1).to(device)
data = np.arange(320).reshape(batch,32,10)
count = np.zeros(batch,256,256)
indices shape is (batch, 10, 2)
</code></pre>
<p>I have spent enormous amount of time trying to optimize this solution, but couldnt. Help is greatly appreciated.<br>
Thanks in advance.</p>

username: ptrblck
cooked: <p>I’m sure there are better approaches and you could speed it up more, but you could use <code>index_put_</code> with the <code>accumulate=True</code> argument instead of the for loop and precompute some indices.<br>
Here is a code snippet:</p>
<pre><code class="lang-python">def fun1():
    result = torch.full(( 32, 256, 256), -1.)
    data = torch.full(( 32, 40000), 1.)
    count = torch.zeros((256,256))
    indices = torch.tensor(([[1,2],[0,0],[1,2],[1,1],[1,0],[0,0],[1,2],[2,2],[2,2],[0,0]]))
    
    for i, indx in enumerate(indices):
            if(count[indx [0], indx [1]] == 0):
                result[:, indx [0], indx [1]] = data [:, i]
            else:
                result[:, indx [0], indx [1]] += data [:, i]
            count[indx [0], indx [1]] += 1
    
    #mask_zero = (count == 0)
    #count[mask_zero] = 1
    #count = torch.Tensor(count)
    #average= result/count
    uidx, ucount = indices.unique(dim=0, return_counts=True)
    result[:, uidx[:, 0], uidx[:, 1]] /= ucount
    return result#average


def fun2():
    result2 = torch.full((32, 256, 256), -1.)
    data = torch.full((32, 40000), 1.)
    indices = torch.tensor([[1,2],[0,0],[1,2],[1,1],[1,0],[0,0],[1,2],[2,2],[2,2],[0,0]])
    
    # prime result2
    result2[:, indices[:, 0], indices[:, 1]] = 0
    
    # precompute fast_idx
    fast_idx = torch.arange(32).unsqueeze(1).expand(-1, 10).transpose(1, 0).reshape(-1)
    
    result2.index_put_(
        (fast_idx,
         indices[:, 0].unsqueeze(1).expand(-1, 32).reshape(-1),
         indices[:, 1].unsqueeze(1).expand(-1, 32).reshape(-1)),
        data[fast_idx, torch.arange(10).unsqueeze(1).expand(-1, 32).reshape(-1)],
        accumulate=True)
    
    uidx, ucount = indices.unique(dim=0, return_counts=True)
    result2[:, uidx[:, 0], uidx[:, 1]] /= ucount
    return result2

avg1 = fun1()
avg2 = fun2()
print((avg1 - avg2).abs().max())
&gt; tensor(0.)

%timeit fun1()
&gt; 6.32 ms ± 754 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit fun2()
&gt; 5.75 ms ± 134 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
</code></pre>
<p>which could potentially speed up your method.<br>
Note that I removed the slower <code>count</code> calculation from your code and replaced it with the <code>unique</code> method, which yields a speedup on my setup.</p>

username: Yasser_Khalil
cooked: <p>Thank you for your help! Indeed I gained good speedup.</p>
<p>What if we level up the problem to include multiple batches (say 8) instead of just one. Noting that each batch can have different number of effective indicies. the ineffictive indicies can be filled with dummy data like (-1). Example</p>
<pre><code class="lang-auto">
indices_0 = torch.tensor(([[1,2],[0,0],[1,2],[1,1],[1,0],[0,0],[1,2],[2,2],[2,2],[0,0]]))
indices_1 = torch.tensor(([[1,0],[0,0],[-1,-1],[-1,-1],[1,1],[0,0],[1,0],[0,0],[0,1],[0,0]]))
indices_2 = torch.tensor(([[-1,-1],[0,0],[3,1],[0,0],[0,0],[1,1],[1,1],[1,1],[1,2],[2,0]]))
...
indices_7 = torch.tensor(([[2,2],[0,1],[2,2],[1,1],[1,2],[0,0],[1,0],[1,2],[2,1],[0,1]]))
</code></pre>
<p>Here [-1,-1] is to be ignored. I indicated the ineffective indices with [-1,-1] just to spot them out. I can eliminate them from the tensor, but then I wont be able to put them in one tensor of shape (8,10,2) as they are of different lengths, unless there is a way around this too.</p>
<p>Those are 8 indices each representing a batch. The shape of all tensors will be as follows:<br>
Whenever there is an ineffective (not to be counted) index then its corresponding data value will be [-1,-1,…,-1]. Again the -1’s are there to make all tensors of equal lengths.</p>
<pre><code class="lang-auto">result = torch.full((8, 32, 256, 256), -1).to(device)
data = torch.rand((8,32,40000))
indices shape is: (8, 10, 2)
</code></pre>

username: Yasser_Khalil
cooked: <p>Hi,</p>
<p>Is there a way to use average pooling instead of summing and then taking the average?</p>
<p>Thanks.</p>