How can i solve the Input Size Error with Conv2d?

Question

RuntimeError: Expected 3D (unbatched) or 4D (batched) input to conv2d, but got input of size: [16, 1280] even though my inputs.shape is torch.Size([16, 3, 120, 120])

Hey im new to pytorch, so please dont be too harsh with me.

I am trying to train the following model to classify images into 12 labels:

model =  models.efficientnet_b1(weights='DEFAULT').to(device) # put it to GPU

model.classifier = nn.Sequential(
    nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, stride=1, padding=1),
    nn.ReLU(inplace=True),
    nn.MaxPool2d(kernel_size=2, stride=2),
    nn.Flatten(),
    nn.Linear(32*16*16, 128),
    nn.ReLU(inplace=True),
    nn.Linear(128, 12)
).to(device)

for param in model.parameters():
            param.requires_grad = False


optimizer = torch.optim.AdamW(model.classifier.parameters(), lr=5e-4, weight_decay=0.1) 

and train it like this:

# Define the loss function and the optimizer
criterion = nn.CrossEntropyLoss()

# Define the number of training epochs
num_epochs = 10

# Training loop
for epoch in range(num_epochs):
    # Set the model to train mode
    model.train()

    # Initialize the running loss for this epoch
    running_loss = 0.0
    
    # Iterate over the training data
    for i, data in enumerate(train_loader):
        inputs, labels = data
        inputs, labels = inputs.to(device), labels.to(device)

        # Zero the gradients
        optimizer.zero_grad()

        # Forward pass
  
        print(inputs.shape) // Print: torch.Size([16, 3, 120, 120])

        outputs = model(inputs) //Error
        loss = criterion(outputs, labels)

        # Backward pass and optimization
        loss.backward()
        optimizer.step()

        # Update the running loss
        running_loss += loss.item()

My problem is, that i get the error: RuntimeError: Expected 3D (unbatched) or 4D (batched) input to conv2d, but got input of size: [16, 1280] I dont understand why i get this error, because the inputs size is [16, 3, 120, 120] (Like Printed)

Really appreciate your help!

Answer 1

You are replacing the model.classifier with your own module. The problem is that the original module expects an input of the form = (batch, 1280) and you are replacing it by a module that expects an input of the form (batch, channels, height, width). You can do something like this:

net.classifier = nn.Sequential(
    nn.Dropout()
    nn.Linear(1280, 12)
)

It is just a working example, I don't know what you are trying to do! I suppose you want the output equal to 12