维度超出范围（预计在 [-4, 3] 范围内，但得到 64）

Question

I am new to Pytorch and I've been working on training the MLP model using the MNIST dataset. Basically, I am feeding the model with images and labels as an input and training the dataset on it. I am using CrossEntropyLoss() as a loss function, however I am getting the dimension error whenever I run my model.

IndexError                                Traceback (most recent call last)
<ipython-input-37-04f8cfc1d3b6> in <module>()
     47 
     48         # Forward
---> 49         outputs = model(images)
     50 

5 frames
/usr/local/lib/python3.7/dist-packages/torch/nn/modules/flatten.py in forward(self, input)
     38 
     39     def forward(self, input: Tensor) -> Tensor:
---> 40         return input.flatten(self.start_dim, self.end_dim)
     41 
     42     def extra_repr(self) -> str:

IndexError: Dimension out of range (expected to be in range of [-4, 3], but got 64)

Here is the MLP class that I've created

class MLP(nn.Module):
    def __init__(self, device, input_size = 1*28*28, output_size = 10):
        super().__init__()
        
        self.seq = nn.Sequential(nn.Flatten(BATCH=64, input_size),
                                                       nn.Linear(input_size, 32),
                                                       nn.ReLU(),
                                                       nn.Linear(32, output_size))
        
        self.to(device)
        
    def forward(self, x):
        return self.seq(x)

And rest of the training model is

from tqdm.notebook import tqdm
from datetime import datetime

from torch.utils.tensorboard import SummaryWriter
import torch.optim as optim

exp_name = "MLP version 1"

# log_name = "logs/" + exp_name + f" {datetime.now()}"
# print("Tensorboard logs will be written to:", log_name)
# writer = SummaryWriter(log_name)

criterion = nn.CrossEntropyLoss()
model = MLP(device)

optimizer = torch.optim.Adam(model.parameters(), lr = 0.0001)
num_epochs = 10

for epoch in tqdm(range(num_epochs)):
    epoch_train_loss = 0.0
    epoch_accuracy = 0.0
    
    for data in train_loader:
        images, labels = data
        images, labels = images.to(device), labels.to(device)
        images = images.permute(0, 3, 1, 2)
        
        
        optimizer.zero_grad()
        print("hello")
        
        outputs = model(images)
    
        loss = criterion(outputs, labels)
        epoch_train_loss += loss.item()
        
        loss.backward()
        optimizer.step()
        
        accuracy = compute_accuracy(outputs, labels)
        epoch_accuracy += accuracy
    writer.add_scalar("Loss/training", epoch_train_loss, epoch)
    writer.add_scalar("Accuracy/training", epoch_accuracy / len(train_loader), epoch)
    
    print('epoch: %d loss: %.3f' % (epoch + 1, epoch_train_loss / len(train_loader)))
    print('epoch: %d accuracy: %.3f' % (epoch + 1, epoch_accuracy / len(train_loader)))
    
    epoch_accuracy = 0.0
    # The code below computes the validation results
    for data in val_loader:
        images, labels = data
        images, labels = images.to(device), labels.to(device)
        images = images.permute(0, 3, 1, 2)
        
        model.eval()
        with torch.no_grad():
            outputs = model(images)
            
        accuracy = compute_accuracy(outputs, labels)
        epoch_accuracy += accuracy
    writer.add_scalar("Accuracy/validation", epoch_accuracy / len(val_loader), epoch)
print("finished training")

Any help would be appreciated. Thank you.

Answer 1

nn.Flatten() instead of nn.Flatten(BATCH=64, input_size)

https://pytorch.org/docs/stable/generated/torch.nn.Flatten.html