Class CheckerboardDataModule

data/checkerboarddataset.py:235–267

    return lonlats, labels

class CheckerboardDataModule(pl.LightningDataModule):
    def __init__(self, num_samples=5000, batch_size=1000, num_classes = 4, num_support = 200):
        super().__init__()
        self.num_samples = num_samples
        self.batch_size=batch_size
        self.num_support = num_support
        self.num_classes = num_classes

        # mean and std distance between clusters given the number of points
        self.mean_dist, self.std_dist = calc_avg_distances(num_support, unit="deg")

    def setup(self, stage: str):
        self.train_ds = TensorDataset(*get_data(N_samples = self.num_samples,
                                                N_support = self.num_support,
                                                n_classes=self.num_classes,
                                                seed=0))
        self.valid_ds = TensorDataset(*get_data(N_samples = self.num_samples,
                                                N_support = self.num_support,
                                                n_classes=self.num_classes,
                                                seed=1))
        self.evalu_ds = TensorDataset(*get_data(N_samples = self.num_samples,
                                                N_support = self.num_support,
                                                n_classes=self.num_classes,
                                                grid=True))

    def train_dataloader(self):
        return DataLoader(self.train_ds, batch_size=self.batch_size, shuffle=True)

    def val_dataloader(self):
        return DataLoader(self.valid_ds, batch_size=self.batch_size, shuffle=False)

    def test_dataloader(self):
        return DataLoader(self.evalu_ds, batch_size=self.batch_size, shuffle=False)

if __name__ == '__main__':
    main()