| 267 | |
| 268 | class Inat2018DataModule(pl.LightningDataModule): |
| 269 | def __init__(self, root, batch_size=2000, mode="location", num_workers=0): |
| 270 | super().__init__() |
| 271 | self.root = root |
| 272 | self.batch_size = batch_size |
| 273 | self.mode = mode |
| 274 | self.num_workers = num_workers |
| 275 | |
| 276 | with open(os.path.join(root, "categories.json")) as da: |
| 277 | data = json.load(da) |
| 278 | df = pd.DataFrame(data)[["name", "id"]] |
| 279 | |
| 280 | self.name2id = df.set_index("name").to_dict()["id"] |
| 281 | self.id2name = df.set_index("id").to_dict()["name"] |
| 282 | |
| 283 | |
| 284 | def setup(self, stage=None): |