Class STL10LinearProbeDataModule

datamodule.py:226–263

        )

class STL10LinearProbeDataModule(_BaseDataModule):
    num_classes = 10
    def prepare_data(self):
        pass

    def transforms(self):
        tform = transforms.Compose([
            transforms.Resize(self.size),
            transforms.ToTensor(),
            transforms.Normalize((0.4467, 0.4398, 0.4066), (0.2603, 0.2566, 0.2713))
        ])
        return tform

    def setup(self, stage=None):
        self.data_train = STL10(
            root = self.data_dir,
            split = 'train', 
            transform = self.transforms()
        )
        
        data_val_test = STL10(
            root = self.data_dir,
            split = 'test', 
            transform = self.transforms()
        )
        test_size = int(0.7 * len(data_val_test))
        val_size = len(data_val_test) - test_size
        self.data_val, self.data_test = random_split(data_val_test, [val_size, test_size])

    def test_dataloader(self):
        return DataLoader(
            dataset = self.data_test,
            batch_size = self.batch_size,
            num_workers = self.num_workers,
            pin_memory = self.pin_memory,
            shuffle = False,
            drop_last = True
        )