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 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)
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)
class CfgpTrainer:
def __init__(self, conv_trainer, gpytorch_trainer):
"""
The `conv_trainer` needs to be a `SimpleTrainer` whose model has the
`_convolve` method. The `gpytorch_trainer` needs to be a
`GPyTorchTrainer`.
"""
self.conv_trainer = conv_trainer
self.gpytorch_trainer = gpytorch_trainer
self.device = self.conv_trainer.device
self.train_loader = self.conv_trainer.train_loader
self.val_loader = self.conv_trainer.val_loader
self.test_loader = self.conv_trainer.test_loader
def train(self, lr=0.1, n_training_iter=20):
print("### Beginning training on convolutional network.")
self.conv_trainer.train()
self._train_gp(lr, n_training_iter)
def _train_gp(self, lr, n_training_iter):
print("### Beginning training on GP.")
convolutions, train_y = self._get_training_convolutions()
self.gpytorch_trainer.train(
train_x=convolutions,
train_y=train_y,
lr=lr,
n_training_iter=n_training_iter,
)
def _get_training_convolutions(self):
train_convs, train_y = self._get_convolutions(self.train_loader)
self.conv_normalizer = Normalizer(train_convs, self.device)
normed_convs = self.conv_normalizer.norm(train_convs)
return normed_convs, train_y
def _get_convolutions(self, data_loader):
self.conv_trainer.model.eval()
convolutions = []
targets = []
for i, batch in enumerate(data_loader):
batch.to(self.device)
out = self.conv_trainer.model._convolve(batch)
for conv, target in zip(out.tolist(), batch.y):
convolutions.append(conv)
targets.append(target)
convolutions = torch.Tensor(convolutions).to(self.device)
targets = torch.Tensor(targets).to(self.device)
return convolutions, targets
def predict(self, src, batch_size=32):
print("### Generating predictions on {}.".format(src))
# Parse the data
dataset_config = {"src": src}
dataset = registry.get_dataset_class(
self.conv_trainer.config["task"]["dataset"])(dataset_config)
data_loader = dataset.get_full_dataloader(batch_size=batch_size)
# Get the convolutions
convs, targets_actual = self._get_convolutions(data_loader)
try:
normed_convs = self.conv_normalizer.norm(convs)
except AttributeError as error:
raise type(error)(str(error) + "; error may have occurred "
"because the CFGP may not have been trained yet")
# Feed the convolutions into the GP
targets_pred, targets_std = self.gpytorch_trainer.predict(normed_convs)
return targets_pred, targets_std
def save_state(self,
gp_path='gp_state.pth',
normalizer_path='normalizer.pth'):
self.gpytorch_trainer.save_state(gp_path)
with open(normalizer_path, 'wb') as f:
pickle.dump(self.conv_normalizer.state_dict(), f)
def load_state(self, nn_checkpoint_file, gp_checkpoint_file,
normalizer_checkpoint_file):
self._load_conv(nn_checkpoint_file)
self._load_gp(gp_checkpoint_file)
self._load_normalizer(normalizer_checkpoint_file)
def _load_conv(self, nn_checkpoint_file):
self.conv_trainer.load_state(nn_checkpoint_file)
def _load_gp(self, gp_checkpoint_file):
convolutions, train_y = self._get_training_convolutions()
self.gpytorch_trainer.load_state(gp_checkpoint_file, convolutions,
train_y)
def _load_normalizer(self, normalizer_checkpoint_file):
with open(normalizer_checkpoint_file, 'rb') as f:
normalizer_state_dict = pickle.load(f)
self.conv_normalizer.load_state_dict(normalizer_state_dict)
def _get_training_convolutions(self):
train_convs, train_y = self._get_convolutions(self.train_loader)
self.conv_normalizer = Normalizer(train_convs, self.device)
normed_convs = self.conv_normalizer.norm(train_convs)
return normed_convs, train_y
