Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from gcn_lib.sparse.torch_nn import MLP | |
| from model.model import DeeperGCN | |
| import numpy as np | |
| import logging | |
| class PLANet(torch.nn.Module): | |
| def __init__(self, args,saliency=False): | |
| super(PLANet, self).__init__() | |
| # Args | |
| self.args = args | |
| # Molecule and protein networks | |
| self.molecule_gcn = DeeperGCN(args, saliency=saliency) | |
| self.target_gcn = DeeperGCN(args, is_prot=True) | |
| # Individual modules' final embbeding size | |
| output_molecule = args.hidden_channels | |
| output_protein = args.hidden_channels_prot | |
| # Concatenated embbeding size | |
| Final_output = output_molecule + output_protein | |
| # Overall model's final embbeding size | |
| hidden_channels = args.hidden_channels | |
| # Multiplier | |
| if args.multi_concat: | |
| self.multiplier_prot = torch.nn.Parameter(torch.zeros(hidden_channels)) | |
| self.multiplier_ligand = torch.nn.Parameter(torch.ones(hidden_channels)) | |
| elif self.args.MLP: | |
| # MLP | |
| hidden_channel = 64 | |
| channels_concat = [256, hidden_channel, hidden_channel, 128] | |
| self.concatenation_gcn = MLP(channels_concat, norm=args.norm, last_lin=True) | |
| # breakpoint() | |
| indices = np.diag_indices(hidden_channel) | |
| tensor_linear_layer = torch.zeros(hidden_channel, Final_output) | |
| tensor_linear_layer[indices[0], indices[1]] = 1 | |
| self.concatenation_gcn[0].weight = torch.nn.Parameter(tensor_linear_layer) | |
| self.concatenation_gcn[0].bias = torch.nn.Parameter( | |
| torch.zeros(hidden_channel) | |
| ) | |
| else: | |
| # Concatenation Layer | |
| self.concatenation_gcn = nn.Linear(Final_output, hidden_channels) | |
| indices = np.diag_indices(output_molecule) | |
| tensor_linear_layer = torch.zeros(hidden_channels, Final_output) | |
| tensor_linear_layer[indices[0], indices[1]] = 1 | |
| self.concatenation_gcn.weight = torch.nn.Parameter(tensor_linear_layer) | |
| self.concatenation_gcn.bias = torch.nn.Parameter( | |
| torch.zeros(hidden_channels) | |
| ) | |
| # Classification Layer | |
| num_classes = args.nclasses | |
| self.classification = nn.Linear(hidden_channels, num_classes) | |
| def forward(self, molecule, target): | |
| molecule_features = self.molecule_gcn(molecule) | |
| target_features = self.target_gcn(target) | |
| # Multiplier | |
| if self.args.multi_concat: | |
| All_features = ( | |
| target_features * self.multiplier_prot | |
| + molecule_features * self.multiplier_ligand | |
| ) | |
| else: | |
| # Concatenation of LM and PM modules | |
| All_features = torch.cat((molecule_features, target_features), dim=1) | |
| All_features = self.concatenation_gcn(All_features) | |
| # Classification | |
| classification = self.classification(All_features) | |
| return classification | |
| def print_params(self, epoch=None, final=False): | |
| logging.info("======= Molecule GCN ========") | |
| self.molecule_gcn.print_params(epoch) | |
| logging.info("======= Protein GCN ========") | |
| self.target_gcn.print_params(epoch) | |
| if self.args.multi_concat: | |
| sum_prot_multi = sum(self.multiplier_prot) | |
| sum_lig_multi = sum(self.multiplier_ligand) | |
| logging.info("Sumed prot multi: {}".format(sum_prot_multi)) | |
| logging.info("Sumed lig multi: {}".format(sum_lig_multi)) | |