choderalab / pinot

Probabilistic Inference for NOvel Therapeutics
MIT License
15 stars 2 forks source link

[Work-In-Progress] A different implementation of GVAE #34

Closed dnguyen1196 closed 4 years ago

dnguyen1196 commented 4 years ago

Reference issues: #26

A more modular implementation of GVAE that is more compatible with our Net class. Adapted from https://github.com/zfjsail/gae-pytorch/tree/master/gae

A quick look at the interface:

class GCNModelVAE(nn.Module):
    def __init__(self, input_feat_dim, hidden_dim1, hidden_dim2, dropout):
        self.output_regression = torch.nn.ModuleList([ # One can add more GCN layers
            GraphConvolution(hidden_dim1, hidden_dim2, dropout, act=lambda x: x),
            GraphConvolution(hidden_dim1, hidden_dim2, dropout, act=lambda x: x),
        ])

    def forward(self, x, adj):
        Returns the latent representation of the nodes . So that when Net calls
        representation(h), it invokes this function

    def condition(self, x, adj):
        Similar to Net.condition, returns a distribution of the approximate posterior q(z|x, adj)

    def encode_and_decode(self, x, adj):
        calls forward, and then condition, and then sample from the approximate posterior
        is used during training to optimize the ELBO

A training file in gvae_exp.py and train.py. The current training scheme loops through each molecule (1186 of them in the esol data set). However, most of these molecules make up small graphs. As a result, there is over fitting going on as training loss can be driven to 0 but validation and testing "scores" is low.

Also comes with utility function that performs "partition" of the graph into train/test/validation edges. Note that this implementation of VGAE (and the formulation introduced in the Kipf and Welling 16 paper) was geared towards link prediction. Therefore, various utilties are implemented to facilitate this task.

karalets commented 4 years ago

Great thanks, I will review tomorrow with a bit more time on my hands.

yuanqing-wang commented 4 years ago

Thanks this looks nice! Let me take a look at the details in just a bit.

yuanqing-wang commented 4 years ago

btw all the macOS tests are failing because of RAM issue or something. let's just disable them since it's distracting.