How to efficiently concatenate two tensors during an training phase

samsja: <p>Hi everybody,</p>
<p>I am working on an nn.Module.<br>
The forward method of the module takes two tensors as input. It first apply a CNN to both of theses inputs. It need then to concatenate the first output with all the lines inside the second output.</p>
<p>What I did for now is the following:</p>
<pre><code class="lang-auto">class Model(nn.class):

def __init__(self):
        super().__init__()
        self.layer1 = SomeModule1()
        self.layer2 = SomeModule2()

def forward(self,x : torch.Tensor(n), y : torch.Tensor(m,n)):
       x = self.layer1(x)
       y = self.layer1(y)

       cat = torch.zeros(2*n,m)

       for i in range(y.size(0)):
             cat[i] = torch.cat(x,y[i])

      return self.layer2(cat)
</code></pre>
<p>the code work. However I feel that there is a better way to achieve it. I have got some grad_fn= in the tensor grad. I am note sure what it is exactly but I have the feeling that it is not efficient.</p>
<p>What would be the best efficient solution for the concatenation ?</p>
<p>thanks in advance <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

ksz: <p>Hi, here there are 3 solutions - the code is a little chaotic, but it should be understandable <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"> :</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.nn.functional as F

import time

class Model(nn.Module):
    def __init__(self):
            super().__init__()
            self.layer1 = nn.Linear(10, 20)
            self.layer2 = nn.Linear(40, 30)

    def forwardA(self,x,y):
        x = self.layer1(x)
        y = self.layer1(y)

        cat = torch.zeros((y.shape[0], 20*2))
        for i in range(y.shape[0]):
            cat[i] = torch.cat((x, y[i]))
            
        return self.layer2(cat)
    
    def forwardB(self,x,y):
        x = self.layer1(x)
        y = self.layer1(y)

        results = []
        for i in range(y.shape[0]):
            results.append(torch.cat((x, y[i])))

        cat = torch.stack(results)
        
        return self.layer2(cat)
    
    def forwardC(self,x,y):
        x = self.layer1(x)
        y = self.layer1(y)

        x = x.repeat(y.shape[0], 1)
        cat = torch.cat((x,y), dim=1)

        return self.layer2(cat)  
    


net = Model()

#########
x = torch.rand((10))
y = torch.rand((5, 10))

resultA = net.forwardA(x,y)
resultB = net.forwardB(x,y)
resultC = net.forwardC(x,y)

if resultA.equal(resultB):
    print("ok_AB")

if resultA.equal(resultC):
    print("ok_AC")    
    
if resultB.equal(resultC):
    print("ok_BC")    
#########

start = time.time()
for i in range(1000):
    x = torch.rand((10))
    y = torch.rand((5, 10))
    net.forwardA(x,y)
end = time.time()
print("Method A:", end - start)

start = time.time()
for i in range(1000):
    x = torch.rand((10))
    y = torch.rand((5, 10))
    net.forwardB(x,y)
end = time.time()
print("Method B:", end - start)

start = time.time()
for i in range(1000):
    x = torch.rand((10))
    y = torch.rand((5, 10))
    net.forwardC(x,y)
end = time.time()
print("Method C:", end - start)
</code></pre>
<p>Results:</p>
<pre><code class="lang-auto">ok_AB
ok_AC
ok_BC
Method A: 0.3101060390472412
Method B: 0.2148759365081787
Method C: 0.17612218856811523
</code></pre>

samsja: <p>I like the third solution because it does not involve python loop. It seems like the good solution, as the benchmark suggest :).</p>