Combining torch.view and torch.gather throws OOM

username: manurare
cooked: <p>Hi,</p>
<p>When I try to time the following code, the line combining <code>torch.view</code> and <code>torch.gather</code> keeps allocating memory in each iteration until OOM. I have no clue why this is happening. If I move the line outside the for loop then it works fine, but this is not what I want as I would like to time everything.</p>
<pre><code class="lang-auto">device = "cuda" if torch.cuda.is_available() else "cpu"
# Variables
B, C, H, W = (10, 3, 640, 256)
x = torch.randn((B, C, H, W), device=device)
indices_y_src = torch.randint(H, size=(2, 5), device=device)
indices_x_src = torch.randint(W, size=(2, 1), device=device)
indices_src = (indices_y_src * W + indices_x_src).view(1, 1, -1).expand(B, C, -1)
indices_y_tgt = torch.randint(H, size=(2, 5), device=device)
indices_x_tgt = torch.randint(W, size=(2, 1), device=device)
indices_tgt = (indices_y_tgt * W + indices_x_tgt).view(1, 1, -1).expand(B, C, -1)
weights = torch.nn.Parameter(torch.ones([3, 3, 3, 3])).to(device)

x_out = torch.nn.functional.conv2d(x, weights, padding=1)
times = []
for _ in range(5000):
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    start.record()
    values = x_out.view(B, C, -1).gather(-1, indices_src).view(B, C, 2, 5)  # This causes OOM
    mean_val = torch.mean(values, -1, True).expand(-1, -1, -1, 5).reshape(B, C, -1)
    x_out = torch.scatter(x_out.view(B, C, -1), -1, indices_tgt, mean_val).view(B, C, H, W)
    end.record()
    # Waits for everything to finish running
    torch.cuda.synchronize()
    times.append(start.elapsed_time(end))
print(f"{torch.mean(torch.tensor(times))} +- {torch.std(torch.tensor(times))}")
</code></pre>

username: ptrblck
cooked: <p><code>x_out</code> is created in a differentiable way as the output of <code>F.conv2d</code>, which means it’s attached to a computation graph. Inside the loop you are also using differentiable operations to recompute <code>x_out</code>, which will increase the computation graph and thus also the memory usage.<br>
You could avoid it by detaching the newly computed <code>x_out</code> via:</p>
<pre><code class="lang-python">x_out = torch.scatter(x_out.view(B, C, -1), -1, indices_tgt, mean_val).view(B, C, H, W).detach()
</code></pre>
<p>which shows a constant memory usage, but will disallow to compute the gradients through all iterations back to <code>weights</code>.</p>

username: manurare
cooked: <p>Hey,</p>
<p>Thanks for the answer. What I am doing is equivalent in result to this other example. However this is not throwing OOM. Is it because I am modifying <code>x_out</code> in place?</p>
<pre><code class="lang-auto">
for _ in range(5000):
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    start.record()
    values = x_out[:, :, indices_y_src, indices_x_src]
    mean_val = torch.mean(values, -1, True)
    x_out[:, :, indices_y_tgt, indices_x_tgt] = mean_val
    end.record()
    # Waits for everything to finish running
    torch.cuda.synchronize()
    times.append(start.elapsed_time(end))
print(f"{torch.mean(torch.tensor(times))} +- {torch.std(torch.tensor(times))}")
</code></pre>

username: ptrblck
cooked: <p>This might be the case but are you seeing any “disallowed inplace operations” during the backward pass?</p>

username: manurare
cooked: <p>No. For example:</p>
<pre><code class="lang-auto">
for _ in range(20):
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    start.record()
    values = x_out[:, :, indices_y_src, indices_x_src]
    mean_val = torch.mean(values, -1, True)
    x_out[:, :, indices_y_tgt, indices_x_tgt] = mean_val
    loss = torch.nn.MSELoss()(x_out, x)
    loss.backward(retain_graph=True)
    end.record()
    # Waits for everything to finish running
    torch.cuda.synchronize()
    times.append(start.elapsed_time(end))
print(f"{torch.mean(torch.tensor(times))} +- {torch.std(torch.tensor(times))}")
</code></pre>
<p>Correct me if I am wrong but I think it is because in <code>values = x_out[:, :, indices_y_src, indices_x_src]</code> the RHS is implicitly cloning those memory addresses into a new contiguous location and then assigning that to LHS. Any inplace change of <code>x_out</code> happening afterwards does not affect the computation graph. If I change that line with <code>values = x_out.view(B, C, -1).gather(-1, indices_src).view(B, C, 2, 5)</code> then it throws <code>RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation</code>  as I believe that gather does not internally clone anything.</p>
<p>I am doing all the intricate gather/scatter/view operations just to avoid the <code>aten::index_put_</code> happening in  <code>x_out[:, :, indices_y_tgt, indices_x_tgt] = mean_val</code>. The backward pass for that operation is really slow compared with gather/scatter. I am using this “trick” in a bigger model and it does speed it up by 1.08x. Do you know if there is a workaround of bypassing <code>aten::index_put_</code> which do not require gather/scatter?</p>

username: ptrblck
cooked: <p>Could you post the shapes you have used to profile <code>gather/scatter</code> vs. <code>index_put_</code>, please?</p>

username: manurare
cooked: <p>You mean the shapes of the tensors I am working with?</p>
<pre><code class="lang-auto">x_out.shape = (10, 3, 640, 256)
indices_x_src.shape = indices_x_tgt.shape = (2, 1)
indices_x_tgt.shape = indices_y_tgt.shape = (2, 5)
indices_src.shape = indices_tgt.shape = (10, 3, 10)
</code></pre>
<p>With indexing (<code>index_put_</code> in backward) I get an E2E time of 61.06 ms<br>
With gather/scatter I get an E2E time of 49.21 ms</p>

username: manurare
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> Any updates on this?</p>

username: ptrblck
cooked: <p>No, since I didn’t spend time writing a minimal and executable code snippet as a reference for profiling.</p>