def load_MolecueNet(data_args):
""" Attention the multi-task problems not solved yet """
molecule_net_dataset_names = {name.lower(): name for name in MoleculeNet.names.keys()}
dataset = MoleculeNet(root=data_args.dataset_dir, name=molecule_net_dataset_names[data_args.dataset_name.lower()])
# Chem.PeriodicTable.GetElementSymbol()
dataset.data.x = dataset.data.x.float()
dataset.data.y = dataset.data.y.squeeze().long()
dataset.node_type_dict = None
dataset.node_color = None
return dataset
def load_MolecueNet(dataset_dir, dataset_name, task=None):
""" Attention the multi-task problems not solved yet """
molecule_net_dataset_names = {name.lower(): name for name in MoleculeNet.names.keys()}
dataset = MoleculeNet(root=dataset_dir, name=molecule_net_dataset_names[dataset_name.lower()])
dataset.data.x = dataset.data.x.float()
if task is None:
dataset.data.y = dataset.data.y.squeeze().long()
else:
dataset.data.y = dataset.data.y[task].long()
dataset.node_type_dict = None
dataset.node_color = None
return dataset