Pytorch model with multi batchnorm1d layers get errors during inference - "expected 2D or 3D input (got 1D input)"

My model in pytorch with batchnorm1D is like this:

class Discriminator(nn.Module):
def __init__(self, sequenceLength):
    super(Discriminator,self).__init__()
    self.batchnorm1 = nn.BatchNorm1d(sequenceLength)
    self.batchnorm2 = nn.BatchNorm1d(2*sequenceLength)
    self.linear1 = nn.Linear(sequenceLength, 2*sequenceLength)
    self.conv2 = nn.Conv1d(1, 1,kernel_size=3, stride=1, padding=1)
    self.conv3 = nn.Conv1d(1, 1,kernel_size=3, stride=1, padding=1)
    self.linear4 = nn.Linear(2*sequenceLength, 1)
    self.relu = nn.ReLU(0.01)
    self.sigmoid = nn.Sigmoid()

def forward(self, x):
    out = self.batchnorm1(x)
    out = self.linear1(out)
    out = self.relu(out)
    out = self.batchnorm2(out)
    out = out.unsqueeze(1)
    out = self.conv2(out)
    out = self.sigmoid(out)
    out = self.conv3(out)
    out = self.relu(out)
    out = out.squeeze()
    out = self.batchnorm2(out)
    out = self.linear4(out)
    out = self.sigmoid(out)
    return out

My inference code is like this:

Discriminator = torch.load('disc.pth', map_location=torch.device('cpu'))
Discriminator.eval()
embededSeq = Embedding.EmbedOne('sample data')
embededSeq = torch.tensor(embededSeq).float()
embededSeq = embededSeq.unsqueeze(0)
score = PosDiscriminator(embededSeq).detach().numpy()[0]

And I got the error message: "expected 2D or 3D input (got 1D input)" at the out = self.batchnorm2(out) line in the model. I wonder if is it caused by my previous line out = out.squeeze() or not? However, the training code worked fine, only happen during the inference.

Could you please have a look and show me what's wrong?

Thank you in advance,

Solution

Yes, this issue is caused by out.squeeze. In general you should avoid out.squeeze() without any dimension arguments because this removes all dimensions with a size of 1, resulting in a tensor with an indeterminate number of dimensions. Within a model the number of dimensions expected by each layer is pretty much always fixed so this is likely to cause problems.

Let your input be of shape [batch_size,length]. (As a side note some of these layers (e.g. conv1D) expect a channel dimension and all of these layers are compatible with a channel dimension, so you may want to add a channel dimension before input to the model rather than altering the number of dimensions within the model.). Here I use length as an arbitrary, changing length, not a fixed value.

batchnorm1 output is of 2 dimensions [batch_size,length]
linear1 otuput is of 2 dimensions [batch_size,length]
unsqueeze output is of 3 dimensions [batch_size,1,length]
conv2 and sigmoid outputs are of 3 dimensions [batch_size,1,length]
conv3 and relu outputs are of 3 dimensions [batch_size,1,length]
here you call out.squeeze(). In the training case, batch_size is greater than 1, so dimension 0 is not squeezed, and only dimension 1 is removed, producing a tensor of size [batch_size,length].
In the inference case, your batch_size is 1, so dimension 0 and 1 are BOTH squeezed. The result is a tensor of shape [length].
The next batchnorm layer expects either a 2D or 3D input [batch_size,channels,length] or [batch_size, length]. In the case where you have a batch of 1 (inference), the input violates this expectation, causing the error.

FIX:

You can fix this by specifying the dimension to squeeze: out = out.squeeze(1), which will always result in a tensor of shape [batch_size,length].