Understanding FixRes - Fixing the train-test resolution discrepancy

import pandas as pd

import torch

import torch.nn as nn
import torch.optim as optim

from torchvision import transforms, models
from torch.utils.data import Dataset, DataLoader
from torch.optim import lr_scheduler

from PIL import Image

import copy
import os
import time

# Ignore warnings to keep the output nice and pretty
import warnings
warnings.filterwarnings("ignore")

df = pd.read_csv('../wandb_cassava_train_val_split.csv')
df.head()

class CassavaDataset(Dataset):
    
    def __init__(self, df, root_dir, transform=None):
        self.df = df
        self.root_dir = root_dir
        self.transform=transform
        
    def __len__(self):
        return len(self.df)
    
    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()
        
        # We get the path of the image
        img_path = os.path.join(self.root_dir, self.df.iloc[idx, 0])
        
        # We load the image and label
        image = Image.open(img_path)
        label = self.df.iloc[idx, 1]

        # We create a single "sample" with the image and its label
        sample = {"image": image, "label": label}
        
        # We apply transforms if specified
        if self.transform:
            sample["image"] = self.transform(sample["image"])
        
        return sample

data_transforms = {
        'train': transforms.Compose([
            # RandomResizedCrop of size 224 used during training
            transforms.RandomResizedCrop(224),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
        ]),
        'val': transforms.Compose([
            # CenterCrop of size 224 used during testing
            transforms.CenterCrop(224),
            transforms.ToTensor(),
        ]),
    }

def get_datasets(data_transforms):

    # We create a training dataset using rows in our data file that have `is_val` equal to False
    # using the transforms we have created above for training...
    train_ds = CassavaDataset(df[df.is_val == False].reset_index(drop=True), '../train_images', data_transforms['train'])
    
    # ... and a test dataset with rows that have `is_val` equal to True,
    # using the transforms we have created above for testing.
    val_ds = CassavaDataset(df[df.is_val == True].reset_index(drop=True), '../train_images', data_transforms['val'])

    return {
        "train": train_ds,
        "val": val_ds
    }

datasets = get_datasets(data_transforms)
dataset_sizes = {x: len(datasets[x]) for x in ['train', 'val']}

def get_dataloaders(datasets, bs=64):
    
    # We create dataloaders from the datasets to enable retrieval of data in batches,
    # to shuffle the data at every epoch, and to parallelize data retrieval (using 4 workers in this case)
    return {
        "train": DataLoader(datasets["train"], batch_size=bs, shuffle=True, num_workers=4),
        "val": DataLoader(datasets["val"], batch_size=bs, shuffle=True, num_workers=4)
    }

dataloaders = get_dataloaders(datasets)

def train_loop(model, dataloader, criterion, optimizer, ds_size, fine_tune):
    
    # We will use the `fine_tune` parameter to ensure that we are
    # modifying only the last batch normalization layer.
    if fine_tune:
        model.eval()
        model.layer4[2].bn3.train()
        
    # Otherwise, we put the model in `train` mode
    # so that all batch normalization layers are updated
    else:
        model.train()
    
    # We keep a track of the loss over all batches in an epoch
    # to later calculate the average loss...
    running_loss = 0.0
    
    # ... and we also keep track of how many correct predictions the
    # model makes to later calculate the average accuracy.
    running_corrects = 0
    
    # We then loop through every batch in an epoch
    for sample in dataloader:
        
        # Extract the images...
        inputs = sample["image"].to(device)
        
        # ... and their corresponding labels
        labels = sample["label"].to(device)
        
        # Boilerplate PyTorch training code:
        
        # 1. Reset all gradients through optimizer
        optimizer.zero_grad()
        
        # 2. Feed model inputs and capture outputs
        outputs = model(inputs)
        
        # 3. Calculate class predicted by model
        _, preds = torch.max(outputs, 1)
        
        # 4. Use correct `labels` and model `outputs` to calculate loss
        loss = criterion(outputs, labels)

        # 5. Compute gradients using the calculated loss
        loss.backward()
        
        # 6. Use optimizer to update model parameters
        optimizer.step()
        
        # Update variables keeping track of loss and correct predictions
        running_loss += loss.item() * inputs.size(0)
        running_corrects += torch.sum(preds == labels.data)

        
    # Calculate the average loss and average accuracy for the epoch
    epoch_loss = running_loss / ds_size
    epoch_acc = running_corrects.double() / ds_size
    
    return epoch_loss, epoch_acc
    
            
        
        
def eval_loop(model, dataloader, criterion, ds_size):
    
    # We put the model in evaluation mode
    model.eval()

    running_loss = 0.0
    running_corrects = 0
    
    # The code used in the evaluation loop is very similar to that used
    # in the training loop.
    # The only difference is we use `torch.no_grad()` to ensure we
    # don't compute any gradients while in evaluation mode.
    with torch.no_grad():
    
        for sample in dataloader:
            inputs = sample["image"].to(device)
            labels = sample["label"].to(device)
        
            outputs = model(inputs)
            _, preds = torch.max(outputs, 1)
            loss = criterion(outputs, labels)
        
            running_loss += loss.item() * inputs.size(0)
            running_corrects += torch.sum(preds == labels.data)
        
    epoch_loss = running_loss / ds_size
    epoch_acc = running_corrects.double() / ds_size
    
    return epoch_loss, epoch_acc
    


def train(
        model, train_loader, val_loader,
        epochs, criterion, optimizer,
        fine_tune=False
    ):
    
    train_start = time.time()

    # We save the model weights that give the best accuracy
    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0
    
    for epoch in range(epochs):
        
        epoch_start = time.time()
        
        # We run the training loop...
        train_loss, train_acc = train_loop(
            model, train_loader, criterion,
            optimizer, dataset_sizes["train"],
            fine_tune=fine_tune
        )
        
        # ... and then the validation loop...
        val_loss, val_acc = eval_loop(
            model, val_loader, criterion, dataset_sizes["val"]
        )
        
        # ... and save the model if it has the best validation accuracy so far
        if val_acc > best_acc:
            best_acc = val_acc
            best_model_wts = copy.deepcopy(model.state_dict())
            
        time_elapsed = time.time() - epoch_start
        print(f'Epoch: {epoch+1:02} | Epoch Time: {time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s')
        print(f'\tTrain Loss: {train_loss:.3f} | Train Acc: {train_acc*100:.2f}%')
        print(f'\t Val. Loss: {val_loss:.3f} |  Val. Acc: {val_acc*100:.2f}%')
        
        
    time_elapsed = time.time() - train_start
    print()
    print('Training complete in {:.0f}m {:.0f}s'.format(
        time_elapsed // 60, time_elapsed % 60))
    print('Best val Acc: {:4f}'.format(best_acc))

    # Load the weights that gave the best accuracy
    model.load_state_dict(best_model_wts)
    return model

model = models.resnet50()

# Keep the same number of inputs,
# but change the number of outputs to the number of unique labels in our dataset
model.fc = nn.Linear(model.fc.in_features, df.label.unique().size)

# Use the power of the GPU! 💪
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = model.to(device)

N_EPOCHS=10

criterion = nn.CrossEntropyLoss()
optimizer = optim.AdamW(model.parameters(), weight_decay=0.1)

model = train(model, dataloaders["train"], dataloaders["val"], N_EPOCHS, criterion, optimizer)

test_sizes = [64, 128, 224, 384, 480]

accuracy = []

for size in test_sizes:

    # For each of the test_sizes we
    # 1. First resize the image to 1.25 the size we want during testing
    # 2. Take a CenterCrop of the desired size for testing
    trans = transforms.Compose([
            transforms.Resize(int(1.25 * size)),
            transforms.CenterCrop(size),
            transforms.ToTensor(),
        ])

    # We create a new dataset and dataloader with these transforms...
    val_ds = CassavaDataset(df[df.is_val == True].reset_index(drop=True), '../train_images', trans)
    ds_size = len(val_ds)
    dataloader = DataLoader(val_ds, batch_size=64, shuffle=True, num_workers=4)

    # ... and calculate the model's accuracy
    loss, acc = eval_loop(model, dataloader, criterion, ds_size)
    accuracy.append(acc.item())

import matplotlib.pyplot as plt

max_acc = max(accuracy)
max_res = test_sizes[accuracy.index(max_acc)]

plt.plot(test_sizes, accuracy)
plt.scatter(224, 0.7109)
plt.scatter(max_res, max_acc)
plt.show()

# Freeze all layers apart from the fully connected layer
for name, child in model.named_children():
    if 'fc' not in name:
        for _, params in child.named_parameters():
            params.requires_grad = False

data_transforms = {
        'train': transforms.Compose([
            transforms.RandomResizedCrop(384), # RandomResizedCrop of size 384 during training
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
        ]),
        'val': transforms.Compose([
            transforms.Resize(480), # First resize to 1.25 * 384 = 480
            transforms.CenterCrop(384), # and then extract a CenterCrop of size 384
            transforms.ToTensor(),
        ]),
    }

datasets = get_datasets(data_transforms)
dataset_sizes = {x: len(datasets[x]) for x in ['train', 'val']}
dataloaders = get_dataloaders(datasets)

criterion = nn.CrossEntropyLoss()
optimizer = optim.AdamW(model.parameters(), weight_decay=0.1)

model = train(model, dataloaders["train"], dataloaders["val"], 5, criterion, optimizer, fine_tune=True)

from fastai.vision.all import *

# Define the path where the images are
path = Path('..')

# get_x fetches a single image
def get_x(row):
    return path/'train_images'/row['image_id']

# get_y fetches the label of a single image
def get_y(row):
    return row['label']

# The splitter returns the train/test split based on the
# `is_val` column in our CSV file.
def splitter(df):
    train = df.index[~df['is_val']].tolist()
    valid = df.index[df['is_val']].tolist()
    return train, valid

# `get_dls` creates dataloaders with RandomResizedCrop of the specified size
# and also with all the default augmentations
def get_dls(size=224):
    cassava = DataBlock(
        blocks = (ImageBlock, CategoryBlock),
        get_x=get_x,
        get_y=get_y,
        splitter=splitter,
        item_tfms=[RandomResizedCrop(size, min_scale=0.7)],
        batch_tfms=[*aug_transforms()]
    )
    dls = cassava.dataloaders(df, bs=64)
    return dls

dls = get_dls()
learn = cnn_learner(dls, resnet50, metrics=accuracy)
learn.fine_tune(5)

new_dls = get_dls(size=280)
learn.dls = new_dls
learn.freeze()
learn.fit(5)