Class Inat2018DataModule

data/inat2018_loader.py:268–308

           torch.tensor(users), torch.tensor(dates), torch.tensor(keep_inds), torch.tensor(dataids)

class Inat2018DataModule(pl.LightningDataModule):
    def __init__(self, root, batch_size=2000, mode="location", num_workers=0):
        super().__init__()
        self.root = root
        self.batch_size = batch_size
        self.mode = mode
        self.num_workers = num_workers

        with open(os.path.join(root, "categories.json")) as da:
            data = json.load(da)
        df = pd.DataFrame(data)[["name", "id"]]

        self.name2id = df.set_index("name").to_dict()["id"]
        self.id2name = df.set_index("id").to_dict()["name"]


    def setup(self, stage=None):
        if self.mode == "image" or self.mode == "all":
            return_location = self.mode == "all"
            image_dataset = INAT_ImageDataset(self.root, os.path.join(self.root, "train2018.json"),
                                 is_train=True, return_location=return_location)
            self.train_ds, self.valid_ds = data.random_split(image_dataset, 
                                                             [VAL_SPLIT_RATIO, 1 - VAL_SPLIT_RATIO])
            self.test_ds = INAT_ImageDataset(self.root, os.path.join(self.root, "val2018.json"),
                                 is_train=False, return_location=return_location, 
                                 logits_file=os.path.join(self.root, "val_logits.npy"))
        elif self.mode == "location":
            locations, classes, *_ = load_inat_location_data(self.root, "train2018_locations.json", "train2018.json")
            self.train_ds, self.valid_ds = data.random_split(TensorDataset(locations, classes), 
                                                             [VAL_SPLIT_RATIO, 1 - VAL_SPLIT_RATIO])
            locations, classes, *_ = load_inat_location_data(self.root, "val2018_locations.json", "val2018.json")
            self.test_ds = TensorDataset(locations, classes)
            
    def train_dataloader(self):
        return DataLoader(self.train_ds, batch_size=self.batch_size, shuffle=True, num_workers=self.num_workers)

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