Pytorch fine tuned CNN model giving always the same prediction in training and validation data
I decided to move from TensorFlow to Pytorch and I am with some issues in understanding how it works. I tried to follow This Tutorial which has a very simple example of Feature Extraction from ImegeNet CNNs for a binary classification problem.
In summary, in my code, the network is defined and called as follows
#Function for feature extraction
def set_parameter_requires_grad(model, feature_extracting):
    if feature_extracting:
        for param in model.parameters():
            param.requires_grad = False

# Initialize these variables which will be set in this if statement. Each of these
# variables is model specific.
def initialize_model(model_name, num_classes, feature_extract, use_pretrained):   

    model_ft = None
    input_size = 0

    if model_name == &quot;resnet&quot;:
    &quot;&quot;&quot; Resnet18
    &quot;&quot;&quot;
        if use_pretrained:
            model_ft = models.resnet18(weights=&quot;ResNet18_Weights.IMAGENET1K_V1&quot;)
        else:
            model_ft = models.resnet18(weights=None)

        set_parameter_requires_grad(model_ft, feature_extract)
        num_ftrs = model_ft.fc.in_features
        model_ft.fc = nn.Linear(num_ftrs, num_classes)
        input_size = 224

    return model_ft, input_size


model_ft, input_size = initialize_model(&quot;resnet&quot;, 2, feature_extract = True, use_pretrained = True)

The transforms I do in the data are defined below:
data_transforms = {
'train': transforms.Compose([
    transforms.RandomResizedCrop(input_size),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]),
'val': transforms.Compose([
    transforms.Resize(input_size),
    transforms.CenterCrop(input_size),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]),

}
I call the training this way:
if feature_extract:
    params_to_update = []

    for name, param in model_ft.named_parameters():
        if param.requires_grad:
            params_to_update.append(param)
            print(&quot;\t&quot;, name)
else:
    for name, param in model_ft.named_parameters():
        if param.requires_grad:
            print(&quot;\t&quot;, name)

# Observe that all parameters are being optimized
optimizer_ft = optim.SGD(params_to_update, lr=0.001, momentum=0.9)
criterion = nn.CrossEntropyLoss()
model_ft, hist = train_model(model_ft, dataloaders_dict, criterion, optimizer_ft, 
num_epochs=num_epochs, is_inception=(model_name == &quot;inception&quot;))

The model has 94% validation accuracy. After training, I save and load the model again.
# Save the model (SIMPLEST WAY)
print(&quot;save model the easiest way&quot;)
torch.save(model_ft, &quot;MY_FIRST_TORCH_MODEL&quot;)
print(&quot;Load model the easiest way&quot;)
model = torch.load(&quot;MY_FIRST_TORCH_MODEL&quot;)
print(&quot;Evaluate model&quot;)
model.eval()

So here it goes my problem: I don't know how to process a testing data sample in order to make it be tested by the Pytorch model. I tried to load and test a training/validation sample just to be sure that the same accuracies in the training and test will also remain when I do a manual test. Here is what I did:
for name in os.listdir(&quot;./hymenoptera_data/train/bees/&quot;):
    # testing one example image
    print(&quot;Predicting &quot;+ &quot;./hymenoptera_data/train/bees/&quot;+name)
    img = cv2.imread(&quot;./hymenoptera_data/train/bees/&quot;+name)
    img_to_test= cv2.resize(img, (224,224), interpolation = cv2.INTER_AREA)
    print(img_to_test.shape)
    img_to_test = img_to_test.astype('float32')

    test_x = img_to_test.reshape(1, 3, 224, 224)
    test_x = torch.from_numpy(test_x)

    output = model(test_x.cuda())
    pred = torch.argmax(output, dim=1)

    print(output)
    print(pred)

So, no matter which sample I choose to test, the prediction is always the same
What am I doing wrong here?

Should the model be loaded in a different way?
Is my function to test correct, especially considering that here I am doing feature extraction instead of fine tuning?
Should I pre-process the data using the same transforms.compose function used to train and validate?
What else should be added to test correctly an image given the example above?


Be careful, img_to_test is in the HWC format. You are reshaping the image when you should be transposing its axes from HWC to CHW. You may want to replace the following:
&gt;&gt;&gt; test_x = img_to_test.reshape(1, 3, 224, 224)

With a call to np.transpose or torch.transpose instead:
&gt;&gt;&gt; test_x = image_to_test.transpose(2,0,1).unsqueeze(0)

