simclr_ds_datagen.py

class DSDataGen(Dataset):
    def __init__(self, phase, imgarr,labels,num_classes):
        
        self.phase = phase
        self.num_classes = num_classes
        self.imgarr = imgarr
        self.labels = labels
        
        self.randomcrop = transforms.RandomResizedCrop(32,(0.8,1.0))
    
    def __len__(self):
        return self.imgarr.shape[0]
    
    def __getitem__(self,idx):
        
        x = self.imgarr[idx]
        
        img = torch.from_numpy(x).float()
        label = self.labels[idx]

        if self.phase == 'train':
            img  = self.randomcrop(img)

        img = self.preprocess(img)
        
        return img, label
    
    def on_epoch_end(self):
        idx = random.sample(population = list(range(self.__len__())),k = self.__len__())
        self.imgarr = self.imgarr[idx]
        self.labels = self.labels[idx]
        
    def preprocess(self,frame):
        frame = frame / 255.0
        frame = (frame-MEAN)/STD
        return frame

simclr_ds_datagen_decl.py

dg = DSDataGen('train', trimages, trlabels, num_classes=10)

dl = DataLoader(dg,batch_size = 32, drop_last = True)

vdg = DSDataGen('valid', valimages, vallabels, num_classes=10)

vdl = DataLoader(vdg,batch_size = 32, drop_last = True)