def load_task(self):
print("### Loading dataset: {}".format(self.config["task"]["dataset"]))
self.parallel_collater = ParallelCollater(1)
if self.config["task"]["dataset"] == "single_point_lmdb":
self.train_dataset = registry.get_dataset_class(
self.config["task"]["dataset"])(self.config["dataset"])
self.train_sampler = DistributedSampler(
self.train_dataset,
num_replicas=distutils.get_world_size(),
rank=distutils.get_rank(),
shuffle=True,
)
self.train_loader = DataLoader(
self.train_dataset,
batch_size=self.config["optim"]["batch_size"],
collate_fn=self.parallel_collater,
num_workers=self.config["optim"]["num_workers"],
pin_memory=True,
sampler=self.train_sampler,
)
self.val_loader = self.test_loader = None
self.val_sampler = None
if "val_dataset" in self.config:
self.val_dataset = registry.get_dataset_class(
self.config["task"]["dataset"])(self.config["val_dataset"])
self.val_sampler = DistributedSampler(
self.val_dataset,
num_replicas=distutils.get_world_size(),
rank=distutils.get_rank(),
shuffle=False,
)
self.val_loader = DataLoader(
self.val_dataset,
self.config["optim"].get("eval_batch_size", 64),
collate_fn=self.parallel_collater,
num_workers=self.config["optim"]["num_workers"],
pin_memory=True,
sampler=self.val_sampler,
)
else:
raise NotImplementedError
self.num_targets = 1
# Normalizer for the dataset.
# Compute mean, std of training set labels.
self.normalizers = {}
if self.config["dataset"].get("normalize_labels", True):
if "target_mean" in self.config["dataset"]:
self.normalizers["target"] = Normalizer(
mean=self.config["dataset"]["target_mean"],
std=self.config["dataset"]["target_std"],
device=self.device,
)
else:
raise NotImplementedError
def _init_conv_trainer(self, db_dir, indices):
task = {'dataset': 'gasdb',
'description': ('Regression of DFT calculated binding energes'),
'labels': ['binding energy'],
'metric': 'mae',
'type': 'regression'}
model = {'name': 'cgcnn',
'atom_embedding_size': 64,
'fc_feat_size': 128,
'num_fc_layers': 4,
'num_graph_conv_layers': 6,
'regress_forces': False}
dataset = {'src': db_dir,
'train_size': len(indices),
'val_size': 0,
'test_size': 0}
optimizer = {'batch_size': 64,
'lr_gamma': 0.1,
'lr_initial': 0.001,
'lr_milestones': [25, 45],
'max_epochs': 50, # per hallucination batch
'warmup_epochs': 10,
'warmup_factor': 0.2,
'num_gpus': self.num_gpus}
self.cnn_args = {'task': task,
'model': model,
'dataset': dataset,
'optimizer': optimizer,
'identifier': 'cnn'}
self.cnn_trainer = SimpleTrainer(**self.cnn_args)
self.dataset = Gasdb(self.cnn_args['dataset'])
self.collate_fn = ParallelCollater(self.num_gpus)
def lmdb_to_batch(path):
dataset = SinglePointLmdbDataset({'src': path})
collater = ParallelCollater(1)
loader = DataLoader(
dataset,
batch_size=1,
shuffle=True,
pin_memory=True,
num_workers=8,
collate_fn=collater,
)
data = dataset[0]
batch = collater([data])
return batch
def load_task(self):
assert (self.config["task"]["dataset"] == "single_point_lmdb"
), "EnergyTrainer requires single_point_lmdb dataset"
logging.info(f"Loading dataset: {self.config['task']['dataset']}")
self.parallel_collater = ParallelCollater(
0 if self.cpu else 1,
self.config["model_attributes"].get("otf_graph", False),
)
self.val_loader = self.test_loader = self.train_loader = None
self.val_sampler = self.test_sampler = self.train_sampler = None
if self.config.get("dataset", None):
self.train_dataset = registry.get_dataset_class(
self.config["task"]["dataset"])(self.config["dataset"])
self.train_sampler = self.get_sampler(
self.train_dataset,
self.config["optim"]["batch_size"],
shuffle=True,
)
self.train_loader = self.get_dataloader(
self.train_dataset,
self.train_sampler,
)
if self.config.get("val_dataset", None):
self.val_dataset = registry.get_dataset_class(
self.config["task"]["dataset"])(self.config["val_dataset"])
self.val_sampler = self.get_sampler(
self.val_dataset,
self.config["optim"].get("eval_batch_size",
self.config["optim"]["batch_size"]),
shuffle=False,
)
self.val_loader = self.get_dataloader(
self.val_dataset,
self.val_sampler,
)
if self.config.get("test_dataset", None):
self.test_dataset = registry.get_dataset_class(
self.config["task"]["dataset"])(self.config["test_dataset"])
self.test_sampler = self.get_sampler(
self.test_dataset,
self.config["optim"].get("eval_batch_size",
self.config["optim"]["batch_size"]),
shuffle=False,
)
self.test_loader = self.get_dataloader(
self.test_dataset,
self.test_sampler,
)
self.num_targets = 1
# Normalizer for the dataset.
# Compute mean, std of training set labels.
self.normalizers = {}
if self.normalizer.get("normalize_labels", False):
if "target_mean" in self.normalizer:
self.normalizers["target"] = Normalizer(
mean=self.normalizer["target_mean"],
std=self.normalizer["target_std"],
device=self.device,
)
else:
raise NotImplementedError
def load_task(self):
print("### Loading dataset: {}".format(self.config["task"]["dataset"]))
self.parallel_collater = ParallelCollater(self.config["optim"].get(
"num_gpus", 1))
if self.config["task"]["dataset"] == "trajectory_lmdb":
self.train_dataset = registry.get_dataset_class(
self.config["task"]["dataset"])(self.config["dataset"])
self.train_loader = DataLoader(
self.train_dataset,
batch_size=self.config["optim"]["batch_size"],
shuffle=True,
collate_fn=self.parallel_collater,
num_workers=self.config["optim"]["num_workers"],
)
self.val_loader = self.test_loader = None
if "val_dataset" in self.config:
self.val_dataset = registry.get_dataset_class(
self.config["task"]["dataset"])(self.config["val_dataset"])
self.val_loader = DataLoader(
self.val_dataset,
self.config["optim"].get("eval_batch_size", 64),
shuffle=False,
collate_fn=self.parallel_collater,
num_workers=self.config["optim"]["num_workers"],
)
else:
self.dataset = registry.get_dataset_class(
self.config["task"]["dataset"])(self.config["dataset"])
(
self.train_loader,
self.val_loader,
self.test_loader,
) = self.dataset.get_dataloaders(
batch_size=self.config["optim"]["batch_size"],
collate_fn=self.parallel_collater,
)
self.num_targets = 1
# Normalizer for the dataset.
# Compute mean, std of training set labels.
self.normalizers = {}
if self.config["dataset"].get("normalize_labels", True):
if "target_mean" in self.config["dataset"]:
self.normalizers["target"] = Normalizer(
mean=self.config["dataset"]["target_mean"],
std=self.config["dataset"]["target_std"],
device=self.device,
)
else:
self.normalizers["target"] = Normalizer(
tensor=self.train_loader.dataset.data.y[
self.train_loader.dataset.__indices__],
device=self.device,
)
# If we're computing gradients wrt input, set mean of normalizer to 0 --
# since it is lost when compute dy / dx -- and std to forward target std
if "grad_input" in self.config["task"]:
if self.config["dataset"].get("normalize_labels", True):
if "target_mean" in self.config["dataset"]:
self.normalizers["grad_target"] = Normalizer(
mean=self.config["dataset"]["grad_target_mean"],
std=self.config["dataset"]["grad_target_std"],
device=self.device,
)
else:
self.normalizers["grad_target"] = Normalizer(
tensor=self.train_loader.dataset.data.y[
self.train_loader.dataset.__indices__],
device=self.device,
)
self.normalizers["grad_target"].mean.fill_(0)
if self.is_vis and self.config["task"]["dataset"] != "qm9":
# Plot label distribution.
plots = [
plot_histogram(
self.train_loader.dataset.data.y.tolist(),
xlabel="{}/raw".format(self.config["task"]["labels"][0]),
ylabel="# Examples",
title="Split: train",
),
plot_histogram(
self.val_loader.dataset.data.y.tolist(),
xlabel="{}/raw".format(self.config["task"]["labels"][0]),
ylabel="# Examples",
title="Split: val",
),
plot_histogram(
self.test_loader.dataset.data.y.tolist(),
xlabel="{}/raw".format(self.config["task"]["labels"][0]),
ylabel="# Examples",
title="Split: test",
),
]
self.logger.log_plots(plots)
def load_task(self):
print("### Loading dataset: {}".format(self.config["task"]["dataset"]))
self.parallel_collater = ParallelCollater(
1 if not self.cpu else 0,
self.config["model_attributes"].get("otf_graph", False),
)
if self.config["task"]["dataset"] == "trajectory_lmdb":
self.train_dataset = registry.get_dataset_class(
self.config["task"]["dataset"])(self.config["dataset"])
self.train_loader = DataLoader(
self.train_dataset,
batch_size=self.config["optim"]["batch_size"],
shuffle=True,
collate_fn=self.parallel_collater,
num_workers=self.config["optim"]["num_workers"],
pin_memory=True,
)
self.val_loader = self.test_loader = None
if "val_dataset" in self.config:
self.val_dataset = registry.get_dataset_class(
self.config["task"]["dataset"])(self.config["val_dataset"])
self.val_loader = DataLoader(
self.val_dataset,
self.config["optim"].get("eval_batch_size", 64),
shuffle=False,
collate_fn=self.parallel_collater,
num_workers=self.config["optim"]["num_workers"],
pin_memory=True,
)
if "test_dataset" in self.config:
self.test_dataset = registry.get_dataset_class(
self.config["task"]["dataset"])(
self.config["test_dataset"])
self.test_loader = DataLoader(
self.test_dataset,
self.config["optim"].get("eval_batch_size", 64),
shuffle=False,
collate_fn=self.parallel_collater,
num_workers=self.config["optim"]["num_workers"],
pin_memory=True,
)
if "relax_dataset" in self.config["task"]:
assert os.path.isfile(
self.config["task"]["relax_dataset"]["src"])
self.relax_dataset = registry.get_dataset_class(
"single_point_lmdb")(self.config["task"]["relax_dataset"])
self.relax_sampler = DistributedSampler(
self.relax_dataset,
num_replicas=distutils.get_world_size(),
rank=distutils.get_rank(),
shuffle=False,
)
self.relax_loader = DataLoader(
self.relax_dataset,
batch_size=self.config["optim"].get("eval_batch_size", 64),
collate_fn=self.parallel_collater,
num_workers=self.config["optim"]["num_workers"],
pin_memory=True,
sampler=self.relax_sampler,
)
else:
self.dataset = registry.get_dataset_class(
self.config["task"]["dataset"])(self.config["dataset"])
(
self.train_loader,
self.val_loader,
self.test_loader,
) = self.dataset.get_dataloaders(
batch_size=self.config["optim"]["batch_size"],
collate_fn=self.parallel_collater,
)
self.num_targets = 1
# Normalizer for the dataset.
# Compute mean, std of training set labels.
self.normalizers = {}
if self.config["dataset"].get("normalize_labels", False):
if "target_mean" in self.config["dataset"]:
self.normalizers["target"] = Normalizer(
mean=self.config["dataset"]["target_mean"],
std=self.config["dataset"]["target_std"],
device=self.device,
)
else:
self.normalizers["target"] = Normalizer(
tensor=self.train_loader.dataset.data.y[
self.train_loader.dataset.__indices__],
device=self.device,
)
# If we're computing gradients wrt input, set mean of normalizer to 0 --
# since it is lost when compute dy / dx -- and std to forward target std
if self.config["model_attributes"].get("regress_forces", True):
if self.config["dataset"].get("normalize_labels", False):
if "grad_target_mean" in self.config["dataset"]:
self.normalizers["grad_target"] = Normalizer(
mean=self.config["dataset"]["grad_target_mean"],
std=self.config["dataset"]["grad_target_std"],
device=self.device,
)
else:
self.normalizers["grad_target"] = Normalizer(
tensor=self.train_loader.dataset.data.y[
self.train_loader.dataset.__indices__],
device=self.device,
)
self.normalizers["grad_target"].mean.fill_(0)
if (self.is_vis and self.config["task"]["dataset"] != "qm9"
and distutils.is_master()):
# Plot label distribution.
plots = [
plot_histogram(
self.train_loader.dataset.data.y.tolist(),
xlabel="{}/raw".format(self.config["task"]["labels"][0]),
ylabel="# Examples",
title="Split: train",
),
plot_histogram(
self.val_loader.dataset.data.y.tolist(),
xlabel="{}/raw".format(self.config["task"]["labels"][0]),
ylabel="# Examples",
title="Split: val",
),
plot_histogram(
self.test_loader.dataset.data.y.tolist(),
xlabel="{}/raw".format(self.config["task"]["labels"][0]),
ylabel="# Examples",
title="Split: test",
),
]
self.logger.log_plots(plots)