RuntimeError: only batches of spatial targets supported (3D tensors) but got targets of size: : [32]- Swap

username: Swapnil_Baad
cooked: <p>I am transfer learning pretrained ‘swin_tiny_patch4_window7_224’ model, with input image size 224x224x3</p>
<p>I have modified the models input and output layer:</p>
<pre><code class="lang-auto"># Modify the model's input and output layers

model.head = nn.Linear(model.head.in_features, 11)
</code></pre>
<p>Training code snippet:</p>
<pre><code class="lang-auto"># Training loop
for epoch in range(num_epochs):
    train_correct = 0
    train_total = 0
    train_loss = 0.0
    
    # Training phase
    model.train()
    for images, labels in train_loader:
        images = images.to(device)
        print(images.shape)
        #labels2 = torch.argmax(labels, dim=1)
        labels = labels.to(device)
        print(labels.shape)
        
        optimizer.zero_grad()

        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        train_loss += loss.item()
        train_total += labels.size(0)
        train_correct += (torch.argmax(outputs, dim=1) == labels).sum().item()

    # Calculate training accuracy and loss
    train_accuracy = train_correct / train_total
    train_loss /= len(train_loader)
</code></pre>
<p>Images shape: torch.Size([32, 3, 224, 224])<br>
labels shape: torch.Size([32])<br>
outputs.shape: torch.Size([32, 7, 7, 11])<br>
from above code.</p>

username: ptrblck
cooked: <p>I assume you are using <code>nn.CrossEntropyLoss</code> as your <code>criterion</code>.<br>
If so, it seems you are working on a segmentation use case since your model output is a 4D tensor most likely in the shape <code>[batch_size, nb_classes, height, width]</code>. In this case the target should be a 3D tensor in the shape <code>[batch_size, height, width]</code> containing class indices in the range <code>[0, nb_classes-1]</code>.</p>
<p>If you are not working on a segmentation use case the output shape of your model would be wrong.</p>

username: Swapnil_Baad
cooked: <p>I am working on classification case with 11 classes:<br>
I have defined the input and output:</p>
<pre><code class="lang-auto"># Modify the model's input and output layers

model.head = nn.Linear(model.head.in_features, 11)

# Define loss function and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

</code></pre>
<p>Is there something that I am doing wrong?</p>

username: ptrblck
cooked: <p>I don’t know how the model is defined but would guess you are removing needed layers via:</p>
<pre><code class="lang-python">model.head = nn.Linear(model.head.in_features, 11)
</code></pre>
<p>Check what <code>model.head</code> used before you replaced it with a single linear layer as I think it might have used a flattening operation making sure the output of the model has a shape of <code>[batch_size, nb_classes]</code>.</p>

username: Swapnil_Baad
cooked: <p>Before: modification</p>
<p>(norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)<br>
(head): ClassifierHead(<br>
(global_pool): SelectAdaptivePool2d (pool_type=avg, flatten=Identity())<br>
(drop): Dropout(p=0.0, inplace=False)<br>
(fc): Linear(in_features=768, out_features=1000, bias=True)<br>
(flatten): Identity()<br>
)</p>
<p>After modification:<br>
(norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)<br>
(head): Linear(in_features=768, out_features=11, bias=True)</p>

username: ptrblck
cooked: <p>I would again guess that <code>SelectAdaptivePool2d</code> might be flattening the activation, but based on the <a href="https://github.com/huggingface/pytorch-image-models/blob/2d0dbd17e388953ab81a5c56f80074eff962ea6b/timm/layers/adaptive_avgmax_pool.py#L161">code</a> <code>flatten</code> should be a <code>bool</code> value while your code seems to pass an <code>Identity</code> module to it, so you might need to dig into the code and check the internal shapes to see if the activations are indeed flattened.<br>
If so, this op would be missing from your code and would explain the errors.</p>