Inserting zeroed channels in convolutional feature map

andrea_bragagnolo: <p>Hi,<br>
I’d like to build a layer that given the output of a <code>Conv2d</code> inserts a series of zero channels in the feature map, e.g. if I have a tensor <code>t1</code> of shape (Bx3xHxW), I want to obtain a new tensor <code>t2</code> of shape (Bx5xHxW) inserting two blocks of shape (Bx1xHxW) in particular places of tensor <code>t1</code> defined by a list of indexes.</p>
<p>Here a 1D example of the expected procedure:</p>
<pre><code class="lang-auto">t1 = torch.tensor([1, 2, 3, 4, 5])  # Starting tensor of size 5
target_size = 8
t2 = torch.zeros(target_size)
idxs = torch.tensor([0, 3, 4, 6, 7])  # Indices of where to put the elements of t1 in t2
# Insert elements of t1 in t2 obtaining t2 = [1, 0, 0, 2, 3, 0, 4, 5]
</code></pre>
<p>I found three different ways of doing so with 4D tensors (output of <code>Conv2d</code>) that work on cuda devices:</p>
<pre><code class="lang-auto">x = torch.randn(bs, ch, shape, shape)
idx = torch.tensor(...)

# matmul
eye = torch.eye(x.shape[1]).unsqueeze(0).to(device)
idx = idx[None, :, None].expand(eye.shape)
zeros = torch.zeros(x.shape[1], target_ch).unsqueeze(0).to(device)
target = torch.scatter(zeros.permute(0, 2, 1), 1, idx, eye).permute(0, 2, 1)

starter.record()
x_flat = x.view(x.shape[0], x.shape[1], -1)  # inference
expanded_x = torch.matmul(target.permute(0, 2, 1), x_flat).view(x.shape[0], -1, x.shape[2], x.shape[3])  
ender.record()
torch.cuda.synchronize()

#scatter
zeros = torch.zeros(x.shape[0], target_ch, *x.shape[2:]).to(device)
idx = idx[None, :, None, None].expand(x.shape)

starter.record()
expanded_x = torch.scatter(zeros, 1, idx, x)
ender.record()
torch.cuda.synchronize()

#index_select
idxs = []
current = 0
for i in range(target_ch):
    if i in idx:
        idxs.append(current)
        current += 1
    else:
        idxs.append(x.shape[1] - 1)
idxs = torch.tensor(idxs, device=x.device)

starter.record()
x = pad(x, (0, 0, 0, 0, 0, 1))
expanded_x = torch.index_select(x, 1, idxs)
ender.record()
torch.cuda.synchronize()
</code></pre>
<p>My problem is that all of theese methods have a computational time that drastically increase with the batch size.<br>
Is there, by chance, another way to do so that has a lower computational cost?</p>
<p>I attach a plot of batch size vs time for all three procedures.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/5/5/550b81a6ff11ea7fdcfea2be4e4aa4776be4155d.png" data-download-href="https://discuss.pytorch.org/uploads/default/550b81a6ff11ea7fdcfea2be4e4aa4776be4155d" title="plots"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/5/5/550b81a6ff11ea7fdcfea2be4e4aa4776be4155d_2_666x500.png" alt="plots" data-base62-sha1="c8leILI4YPODsydsftQrTjJRatD" width="666" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/5/5/550b81a6ff11ea7fdcfea2be4e4aa4776be4155d_2_666x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/5/5/550b81a6ff11ea7fdcfea2be4e4aa4776be4155d_2_999x750.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/5/5/550b81a6ff11ea7fdcfea2be4e4aa4776be4155d_2_1332x1000.png 2x" data-dominant-color="F7F6F3"><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">plots</span><span class="informations">1920×1440 164 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>