| from os.path import join |
|
|
| import torch |
| from pytorch_lightning import LightningDataModule |
| from pytorch_lightning.utilities import rank_zero_only, rank_zero_warn |
| from torch.utils.data import Subset |
| from torch_geometric.loader import DataLoader |
| from torch_scatter import scatter |
| from tqdm import tqdm |
|
|
| from visnet.datasets import * |
| from visnet.utils import MissingLabelException, make_splits |
|
|
|
|
| class DataModule(LightningDataModule): |
| def __init__(self, hparams): |
| super(DataModule, self).__init__() |
| self.hparams.update(hparams.__dict__) if hasattr(hparams, "__dict__") else self.hparams.update(hparams) |
| self._mean, self._std = None, None |
| self._saved_dataloaders = dict() |
| self.dataset = None |
|
|
| def prepare_dataset(self): |
| |
| assert hasattr(self, f"_prepare_{self.hparams['dataset']}_dataset"), f"Dataset {self.hparams['dataset']} not defined" |
| dataset_factory = lambda t: getattr(self, f"_prepare_{t}_dataset")() |
| self.idx_train, self.idx_val, self.idx_test = dataset_factory(self.hparams["dataset"]) |
| |
| print(f"train {len(self.idx_train)}, val {len(self.idx_val)}, test {len(self.idx_test)}") |
| self.train_dataset = Subset(self.dataset, self.idx_train) |
| self.val_dataset = Subset(self.dataset, self.idx_val) |
| self.test_dataset = Subset(self.dataset, self.idx_test) |
|
|
| if self.hparams["standardize"]: |
| self._standardize() |
|
|
| def train_dataloader(self): |
| return self._get_dataloader(self.train_dataset, "train") |
|
|
| def val_dataloader(self): |
| loaders = [self._get_dataloader(self.val_dataset, "val")] |
| delta = 1 if self.hparams['reload'] == 1 else 2 |
| if ( |
| len(self.test_dataset) > 0 |
| and (self.trainer.current_epoch + delta) % self.hparams["test_interval"] == 0 |
| ): |
| loaders.append(self._get_dataloader(self.test_dataset, "test")) |
| return loaders |
|
|
| def test_dataloader(self): |
| return self._get_dataloader(self.test_dataset, "test") |
|
|
| @property |
| def atomref(self): |
| if hasattr(self.dataset, "get_atomref"): |
| return self.dataset.get_atomref() |
| return None |
|
|
| @property |
| def mean(self): |
| return self._mean |
|
|
| @property |
| def std(self): |
| return self._std |
|
|
| def _get_dataloader(self, dataset, stage, store_dataloader=True): |
| store_dataloader = (store_dataloader and not self.hparams["reload"]) |
| if stage in self._saved_dataloaders and store_dataloader: |
| return self._saved_dataloaders[stage] |
|
|
| if stage == "train": |
| batch_size = self.hparams["batch_size"] |
| shuffle = True |
| elif stage in ["val", "test"]: |
| batch_size = self.hparams["inference_batch_size"] |
| shuffle = False |
|
|
| dl = DataLoader( |
| dataset=dataset, |
| batch_size=batch_size, |
| shuffle=shuffle, |
| num_workers=self.hparams["num_workers"], |
| pin_memory=True, |
| ) |
|
|
| if store_dataloader: |
| self._saved_dataloaders[stage] = dl |
| return dl |
| |
| @rank_zero_only |
| def _standardize(self): |
| def get_label(batch, atomref): |
| if batch.y is None: |
| raise MissingLabelException() |
|
|
| if atomref is None: |
| return batch.y.clone() |
|
|
| atomref_energy = scatter(atomref[batch.z], batch.batch, dim=0) |
| return (batch.y.squeeze() - atomref_energy.squeeze()).clone() |
|
|
| data = tqdm( |
| self._get_dataloader(self.train_dataset, "val", store_dataloader=False), |
| desc="computing mean and std", |
| ) |
| try: |
| atomref = self.atomref if self.hparams["prior_model"] == "Atomref" else None |
| ys = torch.cat([get_label(batch, atomref) for batch in data]) |
| except MissingLabelException: |
| rank_zero_warn( |
| "Standardize is true but failed to compute dataset mean and " |
| "standard deviation. Maybe the dataset only contains forces." |
| ) |
| return None |
|
|
| self._mean = ys.mean(dim=0) |
| self._std = ys.std(dim=0) |
| |
| def _prepare_Chignolin_dataset(self): |
| |
| self.dataset = Chignolin(root=self.hparams["dataset_root"]) |
| train_size = self.hparams["train_size"] |
| val_size = self.hparams["val_size"] |
| |
| idx_train, idx_val, idx_test = make_splits( |
| len(self.dataset), |
| train_size, |
| val_size, |
| None, |
| self.hparams["seed"], |
| join(self.hparams["log_dir"], "splits.npz"), |
| self.hparams["splits"], |
| ) |
|
|
| return idx_train, idx_val, idx_test |
| |
| def _prepare_MD17_dataset(self): |
| |
| self.dataset = MD17(root=self.hparams["dataset_root"], dataset_arg=self.hparams["dataset_arg"]) |
| train_size = self.hparams["train_size"] |
| val_size = self.hparams["val_size"] |
| |
| idx_train, idx_val, idx_test = make_splits( |
| len(self.dataset), |
| train_size, |
| val_size, |
| None, |
| self.hparams["seed"], |
| join(self.hparams["log_dir"], "splits.npz"), |
| self.hparams["splits"], |
| ) |
|
|
| return idx_train, idx_val, idx_test |
|
|
| def _prepare_MD22_dataset(self): |
| |
| self.dataset = MD22(root=self.hparams["dataset_root"], dataset_arg=self.hparams["dataset_arg"]) |
| train_val_size = self.dataset.molecule_splits[self.hparams["dataset_arg"]] |
| train_size = round(train_val_size * 0.95) |
| val_size = train_val_size - train_size |
| |
| idx_train, idx_val, idx_test = make_splits( |
| len(self.dataset), |
| train_size, |
| val_size, |
| None, |
| self.hparams["seed"], |
| join(self.hparams["log_dir"], "splits.npz"), |
| self.hparams["splits"], |
| ) |
|
|
| return idx_train, idx_val, idx_test |
|
|
| def _prepare_Molecule3D_dataset(self): |
| |
| self.dataset = Molecule3D(root=self.hparams["dataset_root"]) |
| split_dict = self.dataset.get_idx_split(self.hparams['split_mode']) |
| idx_train = split_dict['train'] |
| idx_val = split_dict['valid'] |
| idx_test = split_dict['test'] |
| |
| return idx_train, idx_val, idx_test |
| |
| def _prepare_QM9_dataset(self): |
| |
| self.dataset = QM9(root=self.hparams["dataset_root"], dataset_arg=self.hparams["dataset_arg"]) |
| train_size = self.hparams["train_size"] |
| val_size = self.hparams["val_size"] |
| |
| idx_train, idx_val, idx_test = make_splits( |
| len(self.dataset), |
| train_size, |
| val_size, |
| None, |
| self.hparams["seed"], |
| join(self.hparams["log_dir"], "splits.npz"), |
| self.hparams["splits"], |
| ) |
|
|
| return idx_train, idx_val, idx_test |
| |
| def _prepare_rMD17_dataset(self): |
| |
| self.dataset = rMD17(root=self.hparams["dataset_root"], dataset_arg=self.hparams["dataset_arg"]) |
| train_size = self.hparams["train_size"] |
| val_size = self.hparams["val_size"] |
| |
| idx_train, idx_val, idx_test = make_splits( |
| len(self.dataset), |
| train_size, |
| val_size, |
| None, |
| self.hparams["seed"], |
| join(self.hparams["log_dir"], "splits.npz"), |
| self.hparams["splits"], |
| ) |
|
|
| return idx_train, idx_val, idx_test |
| |
