I am tying to write a NN in pytorch that learns to give for example: very simple A+B for two images. Writing it, I got this error. my batch_size is 4.
class myNet(nn.Module):
def __init__(self):
super(myNet, self).__init__()
self.fc1 = nn.Linear(3072, 3072) # set up FC layer
self.fc2 = nn.Linear(3072, 3072) # set up the other FC layer
self.fc3 = nn.Linear(3072, 3072) # set up the other FC layer
def forward(self, input1, input2):
a = self.fc1(input1)
print(a.size())
b = self.fc2(input2)
print(b.size())
# now we can reshape `c` and `f` to 2D and concat them
combined = [a[i]+b[i] for i in range(3072)]
print(combined.size())
out = self.fc3(combined)
print(out.size())
return out
and then:
num_epochs = 10
losses = []
batch_size= 4
for epoch in range(num_epochs):
for i, (im1, im2, mid) in enumerate(trainloader):
#im1= im1/255
im1 = torch.from_numpy(np.array(im1, dtype='float32'))
im1 = im1.to(device)
#im2= im2/255
im2 = torch.from_numpy(np.array(im2, dtype='float32'))
im2 = im2.to(device)
#mid= mid/255
mid = torch.from_numpy(np.array(mid, dtype='float32'))
mid = mid.to(device)
# forwad pass
outputs = model(im1, im2)
#outputs = outputs.double()
# loss
loss = criterion(outputs, mid)
losses.append(loss.item())
# backward pass
optimizer.zero_grad()
loss.backward()
# update parameters
optimizer.step()
# report
if (i + 1) % 50 == 0:
print('Epoch [%2d/%2d], Step [%3d/%3d], Loss: %.4f'
% (epoch + 1, num_epochs, i + 1, train_size // batch_size, loss.item()))
and running it, I get the mentioned error. I know its relate to batches. But do not know how to fix it Any ideas to fix it? or maybe another way to write this code is helpful.
On the following line:
combined = [a[i]+b[i] for i in range(3072)]
You are looping over the first dimension of a
and b
on [0, 3072[
. Given that both tensors are shaped batch first, per PyTorch shape convention, you are going out of bound. What you are probably looking to do is loop over the second dimension instead which would correspond to your feature dimension.
combined = [a[:,i]+b[:,i] for i in range(3072)]
This is identical to using the torch.Tensor.__add__
operator:
combined = a + b
Do note however that this is more of a addition fusion rather than a concatenation fusion. Depending on your needs though, you might have to modify this piece of code with respect to how you actually want the two feature tensors to be merged.