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