Reduce memory footprint when processing same input for multiple linear layers

username: gengala
cooked: <p>Consider the following:</p>
<pre><code class="lang-auto">layer = nn.Conv1d(64, 64, kernel_size=1, groups=4)
</code></pre>
<p>This layer represents 4 linear layers that run in parallel, where each of them processes vectors of 16 dimensions and returns vectors of the same dimensionality.</p>
<p>Now, suppose that we want each of this layers to process the same input <code>x = torch.randn(N, 16)</code>.<br>
To do that, we need to transpose and repeat the input as such:</p>
<pre><code class="lang-auto">x = x.t().repeat(4, 1)
</code></pre>
<p>The input will now have shape <code>(64, N)</code>, and we could run <code>layer(x)</code>.</p>
<p>For very large <code>N</code> or vector dimensions, this repeat operation would allocate a huge tensor which we know represents a single tensor repeated multiple times.<br>
My question is: Can we avoid this memory footprint?</p>
<p>I currently don’t see a way to use a combination of <code>expand</code> and <code>view</code> to avoid this issue.</p>

username: ptrblck
cooked: <p>Since you are just repeating the input a plain <code>nn.Conv1d</code> using <code>in_channels=16</code> would yield the same results as your grouped conv:</p>
<pre><code class="lang-python">layer = nn.Conv1d(64, 64, kernel_size=1, groups=4, bias=False)

x = torch.randn(1, 16, 1)
x_ = x.repeat(1, 4, 1)
out_ = layer(x_)

conv = nn.Conv1d(16, 64, 1, bias=False)
conv.weight.shape
with torch.no_grad():
    conv.weight.copy_(layer.weight)

out = conv(x)

print((out_ - out).abs().max())
# tensor(1.1921e-07, grad_fn=&lt;MaxBackward1&gt;)
</code></pre>

username: gengala
cooked: <p>Thanks!</p>
<p>Apparently, we could also just change the <code>groups</code> attribute in layer as <code>layer.groups=1</code> and then run <code>layer(x)</code>.</p>
<p>Would you see any issue in doing that?</p>

username: ptrblck
cooked: <p>Yes, I don’t believe this would directly work since you would increase the weight matrix from <code>[64, 16, 1]</code> to <code>[64, 64, 1]</code>. Of course you could then try to <code>repeat</code> the <code>16</code> filters again to create a kernel filter with the same weights, but this approach sounds quite wasteful since you would be:</p>
<ul>
<li>repeating the filter kernels</li>
<li>repeating the input tensor</li>
<li>increasing the computational workload</li>
</ul>
<p>without a necessity since <code>nn.Conv1d(16, 64, 1, bias=False)</code> will just work without any repetitions.</p>

username: gengala
cooked: <p>My previous reply was not clear enough.</p>
<p>Take a look at the following snippet:</p>
<pre><code class="lang-auto">import torch

layer = torch.nn.Conv1d(64, 64, kernel_size=1, groups=4, bias=True)
print('groups: ', layer.groups, 'weight: ', layer.weight.shape, 'bias: ', layer.bias.shape)

N = 100
x = torch.randn(16, N)
x_repeated = x.repeat(4, 1)

out_repeated = layer(x_repeated)

layer.groups = 1
print('groups: ', layer.groups, 'weight: ', layer.weight.shape, 'bias: ', layer.bias.shape)
out = layer(x)

print(torch.allclose(out, out_repeated, atol=1e-6))
print((out - out_repeated).abs().max())

</code></pre>