|
|
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 |
|
|
|
