def load_model(self):
# Build model
print("### Loading model: {}".format(self.config["model"]))
# TODO(abhshkdz): Eventually move towards computing features on-the-fly
# and remove dependence from `.edge_attr`.
bond_feat_dim = None
if self.config["task"]["dataset"] in [
"trajectory_lmdb",
"single_point_lmdb",
]:
bond_feat_dim = self.config["model_attributes"].get(
"num_gaussians", 50)
else:
raise NotImplementedError
self.model = registry.get_model_class(self.config["model"])(
self.train_loader.dataset[0].x.shape[-1]
if hasattr(self.train_loader.dataset[0], "x")
and self.train_loader.dataset[0].x is not None else None,
bond_feat_dim,
self.num_targets,
**self.config["model_attributes"],
).to(self.device)
print("### Loaded {} with {} parameters.".format(
self.model.__class__.__name__, self.model.num_params))
if self.logger is not None:
self.logger.watch(self.model)
def load_model(self):
# Build model
if distutils.is_master():
logging.info(f"Loading model: {self.config['model']}")
# TODO(abhshkdz): Eventually move towards computing features on-the-fly
# and remove dependence from `.edge_attr`.
bond_feat_dim = None
if self.config["task"]["dataset"] in [
"trajectory_lmdb",
"single_point_lmdb",
]:
bond_feat_dim = self.config["model_attributes"].get(
"num_gaussians", 50
)
else:
raise NotImplementedError
loader = self.train_loader or self.val_loader or self.test_loader
self.model = registry.get_model_class(self.config["model"])(
loader.dataset[0].x.shape[-1]
if loader
and hasattr(loader.dataset[0], "x")
and loader.dataset[0].x is not None
else None,
bond_feat_dim,
self.num_targets,
**self.config["model_attributes"],
).to(self.device)
if distutils.is_master():
logging.info(
f"Loaded {self.model.__class__.__name__} with "
f"{self.model.num_params} parameters."
)
if self.logger is not None:
self.logger.watch(self.model)
self.model = OCPDataParallel(
self.model,
output_device=self.device,
num_gpus=1 if not self.cpu else 0,
)
if distutils.initialized():
self.model = DistributedDataParallel(
self.model, device_ids=[self.device]
)
