Speeding up large tensor operations

username: Stephane_Hatgiskesse
cooked: <p>I have a large 4-dimensional tensor, <code>self.succ_feats</code>, and would frequently like to access elements in this tensor by passing in a list, <code>cords</code>, which corresponds to a set of elements in this tensor. This list of coordinates is also quite large. How can I speed up this operation, as well as any of the subsequent tensor operations?</p>
<p>For references, <code>self.succ_feats</code> is of shape <code>[70, 10, 10, 400]</code>, and <code>cords</code> is of shape <code>[100000,2,2]</code>, meaning that the resulting tensor <code>selected_succ_feats</code> is of shape <code>[100000, 70, 2, 400]</code> which is both extremely slow to process and usually too large to store all at once.</p>
<pre><code class="lang-auto">selected_succ_feats =[self.succ_feats[:,x,y].double() for x,y in cords.long()]
selected_succ_feats = torch.stack(selected_succ_feats) #shape = [100000, 70, 2, 400]
vs = self.linear1(selected_succ_feats) #shape = [100000, 70, 2, 1]
del selected_succ_feats
v_pi_approx = torch.sum(torch.mul(self.softmax(vs),vs),dim = 1) #shape = [100000, 1, 2, 1]
v_pi_approx = torch.squeeze(v_pi_approx)  #shape = [100000, 2]
</code></pre>

username: J_Johnson
cooked: <aside class="quote no-group" data-username="Stephane_Hatgiskesse" data-post="1" data-topic="176005">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/stephane_hatgiskesse/48/58703_2.png" class="avatar"> Stephane_Hatgiskesse:</div>
<blockquote>
<pre><code class="lang-auto">selected_succ_feats =[self.succ_feats[:,x,y].double() for x,y in cords.long()]
selected_succ_feats = torch.stack(selected_succ_feats) #shape = [100000, 70, 2, 400]
vs = self.linear1(selected_succ_feats) #shape = [100000, 70, 2, 1]
del selected_succ_feats
v_pi_approx = torch.sum(torch.mul(self.softmax(vs),vs),dim = 1) #shape = [100000, 1, 2, 1]
v_pi_approx = torch.squeeze(v_pi_approx)  #shape = [100000, 2]
</code></pre>
</blockquote>
</aside>
<p>Welcome to the forums!</p>
<p>List comprehension can be passed directly into a <code>torch.stack([])</code>. Also, you should have your .double() and .long() called before running the list comprehension. Trying to do these inside of a list comprehension will slow things down immensely.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.nn.functional as F
import time

device="cpu"
start_time=time.time()
succ_feats = torch.rand((70, 10, 10, 400), device = device)
coords = torch.randint(0,10, (100000, 2, 2), device = device)
linear1 = nn.Linear(400, 1).to(device)
print("Tensors made", time.time()-start_time)
selected_succ_feats = torch.stack([succ_feats[:,x,y] for x,y in coords]) #shape = [100000, 70, 2, 400]
print("Selection done", time.time()-start_time)
vs = linear1(selected_succ_feats.float()) #shape = [100000, 70, 2, 1]
print("Linear done", time.time()-start_time)
del selected_succ_feats
v_pi_approx = torch.sum(torch.mul(F.softmax(vs),vs),dim = 1) #shape = [100000, 1, 2, 1]
v_pi_approx = torch.squeeze(v_pi_approx)  #shape = [100000, 2]

print("Completed", time.time()-start_time)
</code></pre>
<p>In fact, you might try different data types to see what works best for what you need. I.e. float or half instead of double might be more than sufficient for your purposes and help give you a big speedup with less memory. My guess is you’re running into cached memory, which will slow things down quite a bit with the read/write involved.</p>
<p>The above took under a minute on cpu with 64gb of ram. If I put the succ_feats to double, it started using cached memory which made the time indeterminant.</p>

username: Stephane_Hatgiskesse
cooked: <p>Ok, thanks! Do you think there might be any benefit to trying to use Pytorch’s Dataloader here?</p>

username: J_Johnson
cooked: <p>Yes, if you can batch the operations out, that should eliminate memory overflows. You could use the DataLoader. In that case, just make clear  your definitions <code>__init__</code>, <code>__getitem__</code> and <code>__len__</code> in the Dataset object. For example:</p>
<pre><code class="lang-auto">class CustomDataset(Dataset):
    def __init__(self,  device):
        self.succ_feats = torch.rand((70, 10, 10, 400), device = device)
        self.coords = torch.randint(0,10, (100000, 2, 2), device = device)

    def __len__(self):
        return len(self.coords.size(0))

    def __getitem__(self, idx):
        x, y = self.coords[idx]
        item = self.succ_feats[ : , x, y] 
        return item
</code></pre>