Different predictions after loading model in different jupyter notebook

[shmahaj]

When trying to generate predictions on the test set after loading model from best checkpoint in a fresh python script yields different results compared to the context of python script the model was trained in. The accuracy in the script with training gave a 57% accuracy for 5 epochs while, it was 6% on the different script.

I have also checked the state dictionaries for the loaded model in both contexts, and compared it to the weight values saved in the checkpoint and they are the same. But the predictions for the images are still very different.

Script 1 with Training

import random
import numpy as np
from tqdm import tqdm
import re
import os
import pickle
import pandas as pd
import torch
from torch.nn import functional as F
from torch import nn
from torch.optim import Adam
from torch.utils.data import DataLoader, random_split, Dataset, Sampler
import pytorch_lightning as pl
from pytorch_lightning.core.lightning import LightningModule
from pytorch_lightning.callbacks.early_stopping import EarlyStopping
from pytorch_lightning.callbacks import ModelCheckpoint
from torch.autograd import Variable
import torchvision.models as models
from torchvision import transforms as TF
from torchvision.io import read_image

class RandomSampler(Sampler):
    def __init__(self, data_source, item_dict, num_samples=None):
        self.data_source = data_source
        self._num_samples = num_samples
        self.item_dict = item_dict

        if not isinstance(self.num_samples, int) or self.num_samples <= 0:
            raise ValueError(
                "num_samples should be a positive integer "
                "value, but got num_samples={}".format(self.num_samples)
            )

    @property
    def num_samples(self):
        # dataset size might change at runtime
        if self._num_samples is None:
            return len(self.data_source)
        return self._num_samples

    def __iter__(self):
        n = len(self.data_source)
        # make changes here
        selected_idxs = []
        for l in self.item_dict:
            # randomly select at most 5k
            #print(f"Total samples for the label {l} : {len(self.item_dict[l])}")
            selected = random.sample(self.item_dict[l], min(len(self.item_dict[l]), 5000))
            selected_idxs.extend(selected)
        
        # shuffle the selected_idxs list before returning
        random.shuffle(selected_idxs)
        
        return iter(selected_idxs)
        #return iter(torch.randperm(n, dtype=torch.int64)[: self.num_samples].tolist())

    def __len__(self):
        return self.num_samples

class CustomImageDataset(Dataset):
    def __init__(self, x_train, y_train, label_map, transform=None):
        self.img_labels = y_train
        self.img_data = x_train        
        self.transform = transform
        self.label_map = label_map #{ l:ix for ix,l in enumerate(list(labels))}

    def __len__(self):
        return len(self.img_labels)

    def __getitem__(self, idx):
        img_path = self.img_data.iloc[idx]
        image = read_image(img_path)
        label = self.img_labels.iloc[idx]
        label = self.label_map[label]
        if self.transform:
            image = self.transform(image)
        return image, torch.tensor(label)
    
def get_transform(train):
    transforms = []
    transforms.append(TF.Resize((224,224)))
    transforms.append(TF.ConvertImageDtype(torch.float))
    transforms.append(TF.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)))
    if train:
        transforms.append(TF.RandomHorizontalFlip(0.2))
        transforms.append(TF.RandomVerticalFlip(p=0.2))
        transforms.append(TF.RandomGrayscale(p=0.1))
        transforms.append(TF.RandomAffine(5, translate=None, scale=[1.1,1.1],shear=None, resample=False, fillcolor=0))
        transforms.append(TF.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3, hue=0.3))
        #transforms.append(TF.ToTensor())
        #transforms.append(TF.ConvertImageDtype(torch.float))
        #transforms.append(TF.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)))
    return TF.Compose(transforms)

class ImageClassifier(LightningModule):
    def __init__(self, num_classes, resnet_version, pretrained=False):
        super().__init__()

        self.save_hyperparameters()

        resnets = {
            18: models.resnet18, 34: models.resnet34,
            50: models.resnet50, 101: models.resnet101,
            152: models.resnet152, 51: models.resnext50_32x4d,
            102: models.resnext101_32x8d,
        }

        #instantiate loss criterion
        self.criterion = nn.BCEWithLogitsLoss() if num_classes == 2 else nn.CrossEntropyLoss()
        # Using a pretrained ResNet backbone
        self.resnet_model = resnets[resnet_version](pretrained=pretrained)
        # Replace old FC layer with Identity so we can train our own
        linear_size = list(self.resnet_model.children())[-1].in_features
        # replace final layer for fine tuning
        self.resnet_model.fc = nn.Linear(linear_size, num_classes)

        self.num_classes = num_classes

    def init_datasets(self, train_dset, item_dict, vld_dset, test_dset):
        self.train_dset = train_dset
        self.item_dict = item_dict
        self.test_dset = test_dset
        self.vld_dset = vld_dset

    def configure_optimizers(self):
        return Adam(self.parameters(), lr=1e-3)     

    def forward(self, x):
        # Returns only logits for now

        x = self.resnet_model(x)
        #x = F.log_softmax(x)

        return x

    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        preds = F.log_softmax(logits)

        if self.num_classes == 2:
            y = F.one_hot(y, num_classes=2).float()
        
        loss = self.criterion(logits, y)

        acc = (y == torch.argmax(preds,1)).type(torch.FloatTensor).mean()
        # perform logging
        self.log("train_loss", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True)
        self.log("train_acc", acc, on_step=True, on_epoch=True, prog_bar=True, logger=True)

        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        preds = F.log_softmax(logits)

        if self.num_classes == 2:
            y = F.one_hot(y, num_classes=2).float()
        
        loss = self.criterion(logits, y)
        acc = (y == torch.argmax(preds,1)).type(torch.FloatTensor).mean()
        # perform logging
        self.log("val_loss", loss, on_epoch=True, prog_bar=True, logger=True)
        self.log("val_acc", acc, on_epoch=True, prog_bar=True, logger=True)

    def test_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        preds = F.log_softmax(logits)

        if self.num_classes == 2:
            y = F.one_hot(y, num_classes=2).float()
        
        loss = self.criterion(logits, y)
        acc = (y == torch.argmax(preds,1)).type(torch.FloatTensor).mean()
        # perform logging
        self.log("test_loss", loss, on_step=True, prog_bar=True, logger=True)
        self.log("test_acc", acc, on_step=True, prog_bar=True, logger=True)


    def train_dataloader(self):
        img_train = self.train_dset
        return DataLoader(img_train, batch_size=16, num_workers=6, sampler=RandomSampler(img_train, self.item_dict, 5000))

    def val_dataloader(self):
        img_val = self.vld_dset
        return DataLoader(img_val, batch_size=16, shuffle=False, num_workers=6)

    def test_dataloader(self):
        img_test = self.test_dset
        return DataLoader(img_test, batch_size=16, num_workers=6, shuffle=False)

train_dataset = CustomImageDataset(x_train, y_train, label_map, transform=get_transform(train=True))

val_dataset = CustomImageDataset(x_val, y_val, label_map, transform=get_transform(train=False))

test_dataset =  CustomImageDataset(x_test, y_test, label_map, transform=get_transform(train=False))

model = ImageClassifier(num_classes=n_classes, resnet_version=102, pretrained=True)
model.init_datasets(train_dset=train_dataset, item_dict=item_dict, 
                        vld_dset=val_dataset, test_dset=test_dataset)

trainer_args = {'gpus': 1, 'max_epochs': 5, 'min_epochs': 5}
early_stop_callback = EarlyStopping(monitor="val_loss",  patience=3, verbose=True, mode="min")
MODEL_CKPT_PATH = './resnext101_test3/'


checkpoint_callback = ModelCheckpoint(
    monitor='val_loss',
    dirpath= MODEL_CKPT_PATH,
    filename= 'model-{epoch:02d}-{val_loss:.2f}',
    save_top_k=3,
    mode='min')

trainer = pl.Trainer(**trainer_args, callbacks=[early_stop_callback, checkpoint_callback]) #enable_checkpointing=checkpoint_callback, 
trainer.fit(model)

# Load the best mdoel
model_ckpts = [f for f in os.listdir(MODEL_CKPT_PATH) if f.endswith('.ckpt')]
losses = []
for model_ckpt in model_ckpts:
    loss = re.findall("\d+\.\d+", model_ckpt)
    losses.append(float(loss[0]))

losses = np.array(losses)
best_model_index = np.argsort(losses)[0]
best_model = model_ckpts[best_model_index]
print(best_model)

inference_model = ImageClassifier.load_from_checkpoint(num_classes=n_classes, resnet_version=102, pretrained=False, 
                                                        checkpoint_path=MODEL_CKPT_PATH+best_model)
inference_model.eval()

testloader = DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=6)
preds = []
predicted_labels = []
lls = []
acc = 0
for x in tqdm(testloader):
    #print(x)
    with torch.no_grad():
        outs = inference_model(x[0])
        outs = F.log_softmax(outs, dim=1)
        #print(F.log_softmax(outs))
        for p,l in zip(outs,x[1]):
            #print(p)
            #print(l)
            preds.append(p)
            lls.append(l)
            prediction = torch.argmax(p).item()
            #print(prediction)
            predicted_labels.append(prediction)
            if prediction == l:
                acc += 1

accuracy = acc/len(x_test)
print(accuracy)

The above script yields 57% accuracy.

Script 2 which only loads and tests model

# Load the best mdoel
MODEL_CKPT_PATH = './resnext101_test3/'
model_ckpts = [f for f in os.listdir(MODEL_CKPT_PATH) if f.endswith('.ckpt')]
losses = []
for model_ckpt in model_ckpts:
    loss = re.findall("\d+\.\d+", model_ckpt)
    losses.append(float(loss[0]))

losses = np.array(losses)
best_model_index = np.argsort(losses)[0]
best_model = model_ckpts[best_model_index]
print(best_model)

inference_model = ImageClassifier.load_from_checkpoint(num_classes=n_classes, resnet_version=102, pretrained=False, 
                                                        checkpoint_path=MODEL_CKPT_PATH+best_model)
inference_model.eval()

testloader = DataLoader(test_dataset, batch_size=32, shuffle=False, num_workers=6)
preds = []
predicted_labels = []
lls = []
acc = 0
for x in tqdm(testloader):
    #print(x)
    with torch.no_grad():
        outs = inference_model(x[0])
        outs = F.log_softmax(outs, dim=1)
        #print(F.log_softmax(outs))
        for p,l in zip(outs,x[1]):
            #print(p)
            #print(l)
            preds.append(p)
            lls.append(l)
            prediction = torch.argmax(p).item()
            #print(prediction)
            predicted_labels.append(prediction)
            if prediction == l:
                acc += 1
accuracy = acc/len(x_test)

This only gives an accuracy of 6%.