batch_size isn't match with output_size.size(0)

Question

I am trying to use SelectKBest, mutual_info_classif with CNN model on cifar10 data.

In the for loop,after the model.eval() in the below code, the outputs of model gives (3136,10) tensor size but batch size is 64.

This is the code I run on colab.

 import torch
from torchvision import datasets, transforms
import torch.nn as nn
from torch.utils.data import DataLoader
from sklearn.feature_selection import SelectKBest, mutual_info_classif
import numpy as np

# Define the data transform
data_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# Load the CIFAR-10 dataset
train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=data_transform)
val_dataset = datasets.CIFAR10(root='./data', train=False, download=True, transform=data_transform)

# Create the data loaders for the training and validation sets
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=64, shuffle=False)

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(64*8*8, 128)
        self.fc2 = nn.Linear(128, 10)
    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))
        x = self.pool(torch.relu(self.conv2(x)))
        x = x.view(-1, 64*8*8)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# Initialize the CNN model
model = CNN()

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

# Initialize an empty list to store the features
features = []

model = model.to(device)
# Extract features from the dataset
model.eval()
with torch.no_grad():
    for x, y in train_loader:
        x = x.to(device)
        output = model(x)
        print(output.size())
        output = output.view(output.size(0), -1)
        features.append(output.cpu().numpy())

# Concatenate the features
features = np.concatenate(features)

# Select the k best features using mutual information
selector = SelectKBest(mutual_info_classif, k=1000)
selected_features = selector.fit_transform(features, train_dataset.targets)

Why is this happening?

Answer 1

The size of the output tensor from the model, (3136, 10), suggests that the number of features being extracted from the dataset is 3136, and there are 10 possible classes for the dataset. This is likely the result of the reshaping and flattening of the output from the last convolutional layer in the forward pass of the CNN model, before being passed through the fully connected layers.

In the for loop, the output is being transformed to have size (batch_size, -1) before appending to the features list. This flattens the output to have only 2 dimensions, with the first dimension representing the number of examples in the batch, and the second dimension representing the number of features. Since the batch size is 64, this is likely why the output size is (3136, 10) instead of (64, 10).

The mutual_info_classif function is then used to select the k best features, where k is set to 1000 in this case, from the concatenated features.