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)
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"]))
dataset = registry.get_dataset_class(self.config["task"]["dataset"])(
self.config["dataset"])
if self.config["task"]["dataset"] in ["qm9", "dogss"]:
num_targets = dataset.data.y.shape[-1]
if ("label_index" in self.config["task"]
and self.config["task"]["label_index"] is not False):
dataset.data.y = dataset.data.y[:,
int(self.config["task"]
["label_index"])]
num_targets = 1
else:
num_targets = 1
self.num_targets = num_targets
(
self.train_loader,
self.val_loader,
self.test_loader,
) = dataset.get_dataloaders(
batch_size=int(self.config["optim"]["batch_size"]))
# Normalizer for the dataset.
# Compute mean, std of training set labels.
self.normalizers = {}
if self.config["dataset"].get("normalize_labels", True):
self.normalizers["target"] = Normalizer(
self.train_loader.dataset.data.y[
self.train_loader.dataset.__indices__],
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):
self.normalizers["grad_target"] = Normalizer(
self.train_loader.dataset.data.y[
self.train_loader.dataset.__indices__],
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)
