ValueError: Expected input batch_size (3136) to match target batch_size (64)

LeeCoder: <p>I tried to fix the error, but it didn’t work out, so I’m asking for help. The source code is below.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
import os
import random
import numpy as np
import matplotlib.pyplot as plt
---
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.main = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=10, kernel_size=3,
                      stride=1, padding=1, bias=False),  # 3x32x32 -&gt; 10x32x32
            nn.BatchNorm2d(10),  
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=10, out_channels=20, kernel_size=3,
                      stride=1, padding=1, bias=False),  # 10x32x32 -&gt; 20x32x32
            nn.BatchNorm2d(20),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2),  
            nn.Conv2d(in_channels=20, out_channels=40, kernel_size=3,
                      stride=1, padding=1, bias=False),  
            nn.BatchNorm2d(40),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2), 
            nn.Conv2d(in_channels=40, out_channels=80, kernel_size=3,
                      stride=1, padding=1, bias=False),  
          
            nn.BatchNorm2d(80),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2),  
            nn.Conv2d(in_channels=80, out_channels=100, kernel_size=3,
                      stride=1, padding=1, bias=False),  
            nn.BatchNorm2d(100),
            nn.AdaptiveAvgPool2d(1)  100x1x1
        )
        self.fc = nn.Linear(100, 10)  

    def forward(self, x):
        x = x.float()
        x = x.view(-1, 3, 32, 32)  
        print(x.shape)
        x = self.main(x)  
        x = x.view(-1, 100)  
        print(x.shape)
        x = self.fc(x)  
        return x

print("init model done")
---
batch_size = 64  
test_batch_size = 1000  
max_epochs = 5 
lr = 0.03 
momentum = 0.5 
log_interval = 200  

use_cuda = torch.cuda.is_available()  

device = torch.device("cuda" if use_cuda else "cpu")  

kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}  

print("set vars and device done")
---
transform = transforms.Compose([
                 transforms.Resize((300, 300)),
                 transforms.RandomCrop(224),
                 transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.2),
                 transforms.RandomHorizontalFlip(p = 1),
                 transforms.ToTensor(),  
                 transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])  


train_loader = torch.utils.data.DataLoader(
  datasets.CIFAR10(os.path.join(workspace_path, 'data'), train=True, download=True,
                   transform=transform),
    batch_size = batch_size, shuffle=True, drop_last=True, **kwargs) 

# CIFAR-10 test
test_loader = torch.utils.data.DataLoader(
        datasets.CIFAR10(os.path.join(workspace_path, 'data'), train=False, download=True,
                         transform=transform),
    batch_size=test_batch_size, shuffle=False, **kwargs)
---
model = Net().to(device)  
optimizer = optim.SGD(model.parameters(), lr=lr, momentum=momentum)  
criterion = nn.CrossEntropyLoss()  
---
class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count
---
def train(log_interval, model, device, train_loader, optimizer, epoch):
    model.train()  
    summary_loss = AverageMeter() 
    summary_acc = AverageMeter() 
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device) 
        optimizer.zero_grad()  
        output = model(data)  
        loss = criterion(output, target)  
        loss.backward()  
        optimizer.step()  
        summary_loss.update(loss.detach().item())  
        pred = output.argmax(dim=1, keepdim=True)  
        correct = pred.eq(target.view_as(pred)).sum().item()  
        summary_acc.update(correct / data.size(0))  
        if batch_idx % log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tAverage loss: {:.6f}, Accuracy: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), summary_loss.avg, summary_acc.avg))

    return summary_loss.avg, summary_acc.avg

def test(log_interval, model, device, test_loader):
    model.eval()  
    summary_loss = AverageMeter()  
    summary_acc = AverageMeter() 
    with torch.no_grad(): 
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)  
            output = model(data)  
            loss = criterion(output, target)  
            summary_loss.update(loss.detach().item())  
            pred = output.argmax(dim=1, keepdim=True)  
            correct = pred.eq(target.view_as(pred)).sum().item()  
            summary_acc.update(correct / data.size(0))  

    print('\nTest set: Average loss: {:.4f}, Accuracy: {:.6f}\n'.format
          (summary_loss.avg, summary_acc.avg)) 

    return summary_loss.avg, summary_acc.avg
---
best_acc = 0
best_epoch = 0
for epoch in range(1, max_epochs+1):
    train_loss, train_acc = train(log_interval, model, device, train_loader, optimizer, epoch)
    test_loss, test_acc = test(log_interval, model, device, test_loader)


    if test_acc &gt; best_acc:
        best_acc = test_acc
        best_epoch = epoch
        torch.save(model, os.path.join(workspace_path, f'cifar10_cnn_model_best_acc_{best_epoch}-epoch.pt'))
        print(f'# save model: cifar10_cnn_model_best_acc_{best_epoch}-epoch.pt\n')

print(f'\n\n# Best accuracy model({best_acc * 100:.2f}%): cifar10_cnn_model_best_acc_{best_epoch}-epoch.pt\n')
</code></pre>
<h2><a name="p-456769-torchsize3136-3-32-32-torchsize3136-100-1" class="anchor" href="#p-456769-torchsize3136-3-32-32-torchsize3136-100-1"></a>torch.Size([3136, 3, 32, 32])<br>
torch.Size([3136, 100])</h2>
<p>ValueError                                Traceback (most recent call last)<br>
 in &lt;cell line: 3&gt;()<br>
2 best_epoch = 0<br>
3 for epoch in range(1, max_epochs+1):<br>
----&gt; 4     train_loss, train_acc = train(log_interval, model, device, train_loader, optimizer, epoch)<br>
5     test_loss, test_acc = test(log_interval, model, device, test_loader)<br>
6</p>
<p>4 frames<br>
/usr/local/lib/python3.10/dist-packages/torch/nn/functional.py in cross_entropy(input, target, weight, size_average, ignore_index, reduce, reduction, label_smoothing)<br>
3102     if size_average is not None or reduce is not None:<br>
3103         reduction = _Reduction.legacy_get_string(size_average, reduce)<br>
 → 3104     return torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index, label_smoothing)<br>
3105<br>
3106</p>
<p>ValueError: Expected input batch_size (3136) to match target batch_size (64).</p>

ptrblck: <p>Based on the error message and your code you are changing the <code>batch_size</code> by flattening the tensor in <code>dim0</code> via <code>x = x.view(-1, 3, 32, 32)</code> and <code>x = x.view(-1, 100)</code>. Keep the batch size equal by using <code>x.view(x.size(0), ...)</code> and fix shape mismatches in layers if needed.</p>