How can I properly implement the convolution and summation as shown in the example below?
Lets be given a PyTorch tensor of signals of size (batch_size, num_signals, signal_length)
, i.e. each batch contains several signals. For each batch, I want to convolve the i-th signal with the i-th kernel, and sum all of these convolutions. The result should be of shape (batch_size, 1, signal_length)
The code snipped below should to the right thing, however I want to avoid unnecessary loops.
Thanks for any help!
import torch
import torch.nn as nn
import torch.nn.functional as F
batch_size = 16
num_signals = 5
signal_length = 32
num_kernels = 5
kernel_size = 32
# Random batch of signals and kernels
signals = torch.randn(batch_size, num_signals, signal_length)
kernels = torch.randn(num_kernels, 1, kernel_size)
out = torch.zeros(signals.shape[0], 1, signal_length)
for b in range(signals.shape[0]):
for k in range(num_signals):
out[b, 0, :] += F.conv1d(
signals[b, k].unsqueeze(0).unsqueeze(0),
kernels[k].unsqueeze(0),
padding="same"
).squeeze()
print("Output shape:", out.shape)