hubiodatalab / druggen Goto Github PK

When going through the instructions on the README page, I came across a TypeError after this input: python DrugGEN/main.py --submodel="CrossLoss" --mode="train" --raw_file="DrugGEN/data/chembl_train.smi" --dataset_file="chembl45_train.pt" --drug_raw_file="DrugGEN/data/akt_train.smi" --drug_dataset_file="drugs_train.pt" --max_atom=45

The TypeError was related to the trainer.py file on line 590, specifically this part:
real_graphs, a_tensor, x_tensor = load_molecules(
data=data,
batch_size=self.batch_size,
device=self.device,
b_dim=self.b_dim,
m_dim=self.m_dim,
)

The TypeError was due to load_molecules() not taking 'data' as an argument. This is shown in the training_data.py file:
def load_molecules(batch=None, b_dim=32, m_dim=32, device=None, batch_size=32):
data = data.to(device)
a = geoutils.to_dense_adj(
edge_index = data.edge_index,
batch=data.batch,
edge_attr=data.edge_attr,
max_num_nodes=int(data.batch.shape[0]/batch_size)

After we added in 'data=None', 'after batch_size=32', the training of the model began running.

I've trained a model on the 'Ligand' submodel variation successfully. When running the model to create an output of SMILES I came across this error related to the the tensor shapes:

RuntimeError: Error(s) in loading state_dict for Generator2:
size mismatch for drug_nodes.weight: copying a param with shape torch.Size([128, 9]) from checkpoint, the shape in current model is torch.Size([128, 27]).
size mismatch for nodes_output_layer.weight: copying a param with shape torch.Size([9, 128]) from checkpoint, the shape in current model is torch.Size([27, 128]).
size mismatch for nodes_output_layer.bias: copying a param with shape torch.Size([9]) from checkpoint, the shape in current model is torch.Size([27]).

I can't find anywhere in the code to change the number of nodes to generate molecules

Thank you for your exciting work, which has benefited us a lot.
However, when I was following your work, I wanted to use DrugGEN-Prot (the default model) to train two GAN models at the same time. However, when I train on a single Nvidia 3090 graphics card, there will be a CUDA out of memory situation. However, I see that your running graphics card is similar to mine. How should this situation be solved?