Putting a small tensor to the GPU is adding a huge bottleneck

username: David_Ireland
cooked: <p>I was profiling my code and noticed a list comprehension taking a long time. The list comprehension was in the following method:</p>
<pre><code class="lang-auto">    def get_targets(self, next_states):
        next_states = torch.stack([next_states[:, i, :].unsqueeze(dim=0).repeat(self.ensemble_size, 1, 1) for i in range(self.ensemble_size)])
        with torch.no_grad():
            targets_ = self.critic_target.batched_forward_(next_states).max(dim=-1)[0].mean(dim=-1)
            idx = [torch.randperm(self.ensemble_size)[:2] for _ in range(self.ensemble_size)]
            targets = torch.cat([targets_[i][:, idx[i]].min(dim=-1, keepdim=True)[0] for i in range(self.ensemble_size)], dim=1)
        return targets
</code></pre>
<p>the list comprehension in question is the line which assigns the <code>targets</code> variable. So, I thought I would try and make this quicker by using the following:</p>
<pre><code class="lang-auto">    def get_targets(self, next_states):
        next_states = torch.stack([next_states[:, i, :].unsqueeze(dim=0).repeat(self.ensemble_size, 1, 1) for i in range(self.ensemble_size)])
        with torch.no_grad():
            targets_ = self.critic_target.batched_forward_(next_states).max(dim=-1)[0].mean(dim=-1)
            idx = torch.stack([torch.randperm(self.ensemble_size)[:2] for _ in range(self.ensemble_size)]).to(self.device)
            targets = targets_.gather(-1, idx.unsqueeze(dim=1).repeat(1, self.batch_size, 1)).min(dim=-1)[0].transpose(0, 1)
        return targets
</code></pre>
<p>Now, after profiling, the overall time is slightly faster but there is still a huge bottleneck in putting the <code>idx</code> tensor onto the GPU, and this method is called at the same frequency as the <code>get_targets</code> method. What is most confusing is that I put much larger tensors onto the GPU in other parts of the file and they are much quicker. Does anybody know what could be causing this?</p>
<p>For reproducibility, <code>next_states</code> is a tensor of sizes <code>(512, 25, 9)</code>, replace the line <code>targets_ = self.critic_target.batched_forward_(next_states).max(dim=-1)[0].mean(dim=-1)</code> with <code>targets_ = torch.randn((25, 512, 25))</code>, and <code>self.ensemble_size</code> and <code>self.batch_size</code> are 25 and 512, respectively.</p>
<p>Below, I will also add a snapshot of the profiling results to demonstrate how the <code>.to()</code> method starts to be dominated by this single line. The LHS is the code <em>after</em> making the change. We can see that we get rid of the list comprehension that was dominating the time, but with only 10k more calls to the <code>.to()</code> method (which checks out with how many times the above method was called during profiling) we are now almost 10x slower. As I mentioned previously, I don’t know why this would be the case since I put much larger tensors on the GPU using this code:</p>
<pre><code class="lang-auto">    def get_batch(self):
        batch = self.memory.sample()
        states, actions, rewards, next_states, dones = batch
        states = states.view(self.batch_size, self.ensemble_size, -1).to(self.device)
        actions = actions.view(self.batch_size, self.ensemble_size, -1).to(self.device)
        rewards = rewards.view(self.batch_size, self.ensemble_size).to(self.device)
        next_states = next_states.view(self.batch_size, self.ensemble_size, -1).to(self.device)
        dones = dones.view(self.batch_size, self.ensemble_size).to(self.device)

        return states, actions, rewards, next_states, dones
</code></pre>
<p>where we can assume that the dimensions are much larger than the small <code>idx</code> tensor. In case it is relevant, <code>states</code>, <code>next_states</code>, <code>rewards</code> are all float tensors whilst <code>dones</code> and <code>actions</code> are long tensors.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/8/1/81e0486512abd44b60102532321d1fb56d043fb3.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/81e0486512abd44b60102532321d1fb56d043fb3" title="Screenshot 2023-05-09 at 6.24.55 pm"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/8/1/81e0486512abd44b60102532321d1fb56d043fb3_2_690x104.jpeg" alt="Screenshot 2023-05-09 at 6.24.55 pm" data-base62-sha1="iwW7h6MHa8EVbuoc9WG6UFiLqAb" width="690" height="104" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/8/1/81e0486512abd44b60102532321d1fb56d043fb3_2_690x104.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/8/1/81e0486512abd44b60102532321d1fb56d043fb3_2_1035x156.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/8/1/81e0486512abd44b60102532321d1fb56d043fb3_2_1380x208.jpeg 2x" data-dominant-color="52575A"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot 2023-05-09 at 6.24.55 pm</span><span class="informations">1920×292 48.3 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>