Given
import torch
import torch.nn as nn
nof_samples = 1
nof_features = 2
input = torch.rand(nof_samples, nof_features)
I have two feature networks (one for each feature), which look like this:
feature_networks = nn.ModuleList()
for _ in range(nof_features):
feature_networks.append(
nn.Sequential(
nn.Linear(in_features=1, out_features=3),
nn.Linear(in_features=3, out_features=3),
nn.Linear(in_features=3, out_features=1)
))
output = torch.zeros(nof_samples, nof_features)
for feature_index in range(nof_features): # Iterate over features
output[:, feature_index] = feature_networks[feature_index](input[:, feature_index])
output # shape of [1, 2]
# nof trainable params: 44
I would like to replace it with sth like the following to get rid of the (I believe) time-consuming for-loop:
conv1 = nn.Conv1d(in_channels=nof_features, out_channels=nof_features*3, kernel_size=1, groups=nof_features)
conv2 = nn.Conv1d(in_channels=nof_features, out_channels=nof_features*3, kernel_size=3, groups=nof_features)
conv3 = nn.Conv1d(in_channels=nof_features, out_channels=nof_features, kernel_size=3, groups=nof_features)
outputv2 = conv3(conv2(conv1(input.unsqueeze(2)).view(nof_samples, nof_features, 3)).view(nof_samples, nof_features, 3)).squeeze(2)
outputv2
# shape of [1, 2, 1]
# nof trainable params: 44
Are outputs
and outputsv2
the same (assuming the weights are)? Or do I have to do it differently?