cnn wrong prediction even though model shows good accuracy in training and validation data
I have used the skin cancer classification competition data in Kaggle. There are 4 labels and the entire data is imbalanced. I ran the resnet 18 model on a 10 fold cross validation split to train the data and each fold was given around 2 epochs. The code has been attached below.
Basically the model gave 98.2% accuracy with 0.07 loss value in the train data and 98.1% accuracy and 0.06 loss value in the validation data. So this seemed pretty good.
However the problem is...prediction.py(code attached below). When I tried to predict, the model keeps giving the result as [0]. Even if it's a train image data.
Is there something wrong with my code?
Expected result:
if the image is the input, the output should be either 0,1,2 or 3
model.py(where the training happens)
import pandas as pd 
import numpy as np

import torch
import torch.nn as nn

import os

import cv2

import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split

import albumentations as A

from torch.utils.data import TensorDataset, DataLoader,Dataset
from torchvision import models
from collections import defaultdict
from torch.utils.data.sampler import RandomSampler
import torch.optim as optim
from torch.optim import lr_scheduler
from sklearn import model_selection
from tqdm import tqdm
import gc


# generate data from csv file
class Build_dataset(Dataset):
    def __init__(self, csv, split, mode, transform=None):
        self.csv = csv.reset_index(drop=True)
        self.split = split
        self.mode = mode
        self.transform = transform

    def __len__(self):
        return self.csv.shape[0]

    def __getitem__(self, index):
        row = self.csv.iloc[index]

        image = cv2.imread(row.filepath)
        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)

        if self.transform is not None:
            res = self.transform(image=image)
            image = res['image'].astype(np.float32)
        else:
            image = image.astype(np.float32)

        image = image.transpose(2, 0, 1)
        data = torch.tensor(image).float()

        if self.mode == 'test':
            return data
        else:
            return data, torch.tensor(self.csv.iloc[index].target).long()

# training data           
def train_epoch(model, loader, optimizer,loss_fn,device, scheduler,n_examples):

    model = model.train()

losses = []
correct_predictions = 0

for inputs, labels in tqdm(loader):
    inputs = inputs.to(device)
    labels = labels.to(device)

    outputs = model(inputs)

    _, preds = torch.max(outputs, dim=1)
    loss = loss_fn(outputs, labels)
    
    
    correct_predictions += torch.sum(preds == labels)
    losses.append(loss.item())
    
    

    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
# here you delete inputs and labels and then use gc.collect
    del inputs, labels
    gc.collect()


return correct_predictions.double() / n_examples, np.mean(losses)

# validation data 
def val_epoch(model, loader,loss_fn, device,n_examples):

    model = model.eval()

    losses = []
    correct_predictions = 0

    with torch.no_grad():
        for inputs, labels in tqdm(loader):
            inputs = inputs.to(device)
            labels = labels.to(device)
            outputs = model(inputs)
            _, preds = torch.max(outputs, dim=1)
            loss = loss_fn(outputs, labels)
            correct_predictions += torch.sum(preds == labels)
            losses.append(loss.item())
            # here you delete inputs and labels and then use gc.collect
            del inputs, labels
            gc.collect()
        

    return correct_predictions.double() / n_examples, np.mean(losses)

        


 def train(fold, model,device, num_epochs):

    df_train = df[df.kfold != fold].reset_index(drop=True)
    df_valid = df[df.kfold == fold].reset_index(drop=True)
    # generate data
    dataset_train = Build_dataset(df_train,  'train', 'train', transform=transforms_train)
    dataset_valid = Build_dataset(df_valid, 'train', 'val', transform=transforms_val)

    #load data 
    train_loader = DataLoader(dataset_train, batch_size = 64,sampler=RandomSampler(dataset_train), 
num_workers=4)
    valid_loader = DataLoader(dataset_valid, batch_size = 32,shuffle = True, num_workers= 4 )

    dataset_train_size = len(dataset_train)

    dataset_valid_size = len(dataset_valid)

    optimizer = optim.Adam(model.parameters(), lr = 1e-4)

    model = model.to(device)

    scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, patience = 3,threshold = 0.001, mode = 
'max')

    loss_fn = nn.CrossEntropyLoss().to(device)

    history = defaultdict(list)

    best_accuracy = 0.0

    for epoch in range(num_epochs):
        print(f'Epoch {epoch+1} / {num_epochs}')
        print ('-'*30)
    
        train_acc, train_loss = train_epoch(model, train_loader, optimizer, loss_fn, device, 
scheduler, dataset_train_size)
        print(f'Train loss {train_loss} accuracy {train_acc}')
        valid_acc, valid_loss = val_epoch(model, valid_loader, loss_fn, device,dataset_valid_size)
        print(f'Val   loss {valid_loss} accuracy {valid_acc}')
        print()
    
        history['train_acc'].append(train_acc)
        history['train_loss'].append(train_loss)
        history['val_acc'].append(valid_acc)
        history['val_loss'].append(valid_loss)
    
        if valid_acc &gt; best_accuracy:
            print('saving model')
            torch.save(model.state_dict(), f'best_model_{fold}.bin')
            best_accuracy = valid_acc
    
    print(f'Best Accuracy: {best_accuracy}')

    model.load_state_dict(torch.load(f'best_model_{fold}.bin'))

    return model, history



 if __name__ == '__main__':
    #competition data -2020
    data_dir = &quot;../input/jpeg-melanoma-384x384&quot;
    #competition data - 2019
    data_dir2 = &quot;../input/jpeg-isic2019-384x384&quot;
    # device
    device = torch.device(&quot;cuda&quot;)

    # augmenting images


    image_size = 384
    transforms_train = A.Compose([
        A.Transpose(p=0.5),
        A.VerticalFlip(p=0.5),
        A.HorizontalFlip(p=0.5),
        A.RandomBrightness(limit=0.2, p=0.75),
        A.RandomContrast(limit=0.2, p=0.75),
        A.OneOf([
            A.MedianBlur(blur_limit=5),
            A.GaussianBlur(blur_limit=5),
            A.GaussNoise(var_limit=(5.0, 30.0)),
        ], p=0.7),

        A.OneOf([
            A.OpticalDistortion(distort_limit=1.0),
            A.GridDistortion(num_steps=5, distort_limit=1.),
            A.ElasticTransform(alpha=3),
        ], p=0.7),

        A.CLAHE(clip_limit=4.0, p=0.7),
        A.HueSaturationValue(hue_shift_limit=10, sat_shift_limit=20, val_shift_limit=10, p=0.5),
        A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.1, rotate_limit=15, border_mode=0, p=0.85),
        A.Resize(image_size, image_size),
        A.Cutout(max_h_size=int(image_size * 0.375), max_w_size=int(image_size * 0.375), num_holes=1, 
p=0.7),    
        A.Normalize()
    ])

    transforms_val = A.Compose([
        A.Resize(image_size, image_size),
        A.Normalize()
    ])
    # create data
    df_train = pd.read_csv(os.path.join(data_dir, &quot;train.csv&quot;))  #/kaggle/input/siim-isic-melanoma-classification/train.csv
    df_train.head()

    df_train['is_ext'] = 0
    df_train['filepath'] = df_train['image_name'].apply(lambda x: os.path.join(data_dir, 'train', f'{x}.jpg'))

    # dataset from 2020 data
    df_train['diagnosis'] = df_train['diagnosis'].apply(lambda x: x.replace('seborrheic keratosis', 'BKL'))
    df_train['diagnosis'] = df_train['diagnosis'].apply(lambda x: x.replace('lichenoid keratosis', 'BKL'))
    df_train['diagnosis'] = df_train['diagnosis'].apply(lambda x: x.replace('solar lentigo', 'BKL'))
    df_train['diagnosis'] = df_train['diagnosis'].apply(lambda x: x.replace('lentigo NOS', 'BKL'))
    df_train['diagnosis'] = df_train['diagnosis'].apply(lambda x: x.replace('cafe-au-lait macule', 'unknown'))
    df_train['diagnosis'] = df_train['diagnosis'].apply(lambda x: x.replace('atypical melanocytic proliferation', 'unknown'))

        
    # shuffle data
    df = df_train.sample(frac=1).reset_index(drop=True)

    # creating 8 different target values
    new_target = {d: idx for idx, d in enumerate(sorted(df.diagnosis.unique()))}
    df['target'] = df['diagnosis'].map(new_target)
    mel_idx = new_target['melanoma']

    # creating 10 fold cross validation data
    df = df_train.sample(frac=1).reset_index(drop=True)
    df['kfold'] = -1
    y = df_train.target.values
    kf = model_selection.StratifiedKFold(n_splits=10,shuffle=True)
    idx = kf.get_n_splits(X=df,y=y)
    print(idx)
    for fold,(x,y) in enumerate(kf.split(X=df,y=y)):
        df.loc[y,'kfold'] = fold

    df = df[['filepath','diagnosis', 'target', 'is_ext', 'kfold']]

    class_names = list(df['diagnosis'].unique())


    # create model

    def create_model(n_classes):
        model = models.resnet18(pretrained=True)

        n_features = model.fc.in_features
        model.fc = nn.Linear(n_features, n_classes)
        return model.to(device)
    
    base_model = create_model(len(class_names)) # model ready
    
    
    
    # run the model
    for i in range(10):
        #train
        base_model, history = train(i, base_model, device, num_epochs = 2) # train data

prediction.py
from torchvision import models
import torch 
import torch.nn as nn
import albumentations as A
import cv2
import os 
import numpy as np

device = torch.device(&quot;cuda&quot;)
MODEL = None
MODEL_PATH = &quot;../input/prediction/best_model_4.bin&quot;


def create_model(n_classes):
    model = models.resnet18(pretrained=True)

    n_features = model.fc.in_features
    model.fc = nn.Linear(n_features, n_classes)
    return model.to(device)
# generate the data to tensor with transform application

# converting the image to tensor by using the transforms function

class get_image:
    def __init__(self, image_path, targets, transform = None):
        self.image_path = image_path
        self.targets = targets
        self.transform = transform

    def __len__(self):
        return len(self.image_path)
    def __getitem__(self, item):
        targets = self.targets[item]
        resize = 384
        image = cv2.imread(self.image_path[item])
        image = cv2.resize(image, (resize, resize))
        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    
        if self.transform is not None:
            res = self.transform(image = image)
            image = res['image'].astype(np.float32)
        image = image.transpose(2, 0, 1)
        data = torch.tensor(image).float()
        targets = torch.tensor(targets)
    
        return data, targets
    
# load the data by using torch data
# predict values 

# predict function
def predict(image_path, model, model_path):
    image_size = 384

    transforms_val = A.Compose([
        A.Resize(image_size, image_size),
        A.Normalize()
    ])

    test_images = [image_path]
    test_targets = [0]

    test_data = get_image(
        image_path = test_images,
        targets = test_targets,
        transform=transforms_val)
    # loading the data
    test_dataloader = torch.utils.data.DataLoader(test_data, batch_size=1, shuffle = False, 
num_workers=0)
    model = create_model(n_classes = 4)
    model.load_state_dict(torch.load(model_path))
    model.to(device)
    model.eval()
    prediction = []

    with torch.no_grad():
        for test_data, test_target in test_dataloader:
            test_data = test_data.to(device)
            test_target = test_target.to(device)
        
            outputs = model(test_data)
            _,preds = torch.max(outputs.cpu(), 1)
        
            #prediction.extend(preds)
        
            prediction = np.vstack((preds)).ravel()
        
            return prediction
        
def upload_predict():
    image_file = &quot;../input/whatever/ISIC_0075663.jpg&quot;

    if image_file:
        pred = predict(image_file, MODEL, MODEL_PATH)
        print(pred)
    
    return pred

the label and it's count is given right here
3    27126
2     5193
1      584
0      223

Here 0 is considered malignant type cancer and the other labels are of different types.
Here is the link to the data: https://www.kaggle.com/cdeotte/jpeg-melanoma-384x384

I think you might have the answer to your question! You said:

There are 4 labels and the entire data is imbalanced

Assuming that label 0 is no cancer and 1, 2, 3 are cases with different types of skin cancer. If you said that prediction classes are imbalanced, I'm guessing that 98% of the entire sample is 0, so your algorithm simply predicts every case to be 0 so that it will get right 98% of the time. When your algorithm gets to your test set, it will simply predict everything to be 0.
So the problem isn't with your code. You must balance your dataset by upsampling minority classes, downsampling majority class, assigning a weight/bias to your data or using some sort of model ensemble see https://elitedatascience.com/imbalanced-classes. Check out the credit card fraud detection tutorials such as https://towardsdatascience.com/credit-card-fraud-detection-1b3b3b44109b.
