I have trained two models that use the same sequence of image augmentations but in Torchvision and Kornia and I’m observing a significant difference in the performance of these models. I understand that despite fixing random seeds, these augmentations might still be different which might cause some difference in the test accuracy but on average, I assume that both of these models should end with similar accuracy which is not the case.
# PyTorch transformation
train_orig_transform = transforms.Compose([
    transforms.RandomResizedCrop(32),
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.RandomApply([transforms.ColorJitter(0.4, 0.4, 0.4, 0.1)], p=0.8),
    transforms.RandomGrayscale(p=0.2),
    transforms.ToTensor(),
    transforms.Normalize([0.4914, 0.4822, 0.4465], [0.2023, 0.1994, 0.2010])
    ])

This is the kornia version of the above code

class KorniaAugmentationModule(nn.Module):
    def __init__(self, batch_size=512):
        super().__init__()
        # These are standard values for CIFAR10
        self.mu = torch.Tensor([0.4914, 0.4822, 0.4465])
        self.sigma = torch.Tensor([0.2023, 0.1994, 0.2010])

        self.hor_flip_prob = 0.5
        self.jit_prob = 0.8
        self.gs_prob = 0.2

        self.crop = K.RandomResizedCrop(size=(32, 32), same_on_batch=False)
        self.hor_flip = K.RandomHorizontalFlip(p=self.hor_flip_prob, same_on_batch=False)
        self.jit = K.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.1, p=self.jit_prob, same_on_batch=False)
        self.rand_grayscale =  K.RandomGrayscale(p=self.gs_prob, same_on_batch=False)
        
        self.normalize = K.Normalize(self.mu, self.sigma)

    # Note that I should only normalize in test mode; no other type of augmentation should be performed
    
    def forward(self, x, params=None, mode='train'):
        B = x.shape[0]
        if mode == 'train':
                x = self.crop(x, params['crop_params'])
                x = self.hor_flip(x, params['hor_flip_params'])
                x[params['jit_batch_probs']] = self.jit(x[params['jit_batch_probs']], params['jit_params'])
                x = self.rand_grayscale(x, params['grayscale_params'])

        x = self.normalize(x)
        return x

This is the only difference between the two models. Rest of the code for training and testing is a common shared code.
These are the plots I obtain for training loss and testing accuracy of the models.
Plots for Training Loss curve and Test Accuracy - Kornia (orange) vs Torchvision (green)
Screen Shot 2020-09-27 at 6.56.35 PM1428×562 52.9 KB
The difference in test accuracies between the two models is nearly ~11% which is very significant.
Is this expected? Could you suggest ways to debug this issue?