I have two matrices; I need to do convolution between each column of matrix A and the correspondent column of matrix B and store the results in matrix y. This can be done using a for loop:
for i in range(5):
y[:,i]= F.conv1d(A[:, i], B[:, i])
I am wondering if there is any more efficient way to implement this process? For example, avoid using a for loop.
I don’t know if I understand your use case correctly since you did not post an executable code snippet, but assuming you want to use a “column” of B
as the conv kernel, you could avoid the for loop by using a grouped convolution with a permutation in the kernel dimensions as seen here:
A = torch.randn(1, 10, 7)
B = torch.randn(1, 10, 5)
y = torch.zeros(1, 10, 3)
for i in range(A.size(1)):
y[:, i]= F.conv1d(A[:, i:i+1], B[:, i:i+1])
out = F.conv1d(A, B.permute(1, 0, 2), groups=10)
print((out - y).abs().max())
# tensor(3.5763e-07)
Thanks, that is really helpful.