coprinet/pricePrediction/nets/netsGraph.py

from collections import OrderedDict
from typing import Iterable

import torch
import numpy as np
from torch import nn
import torch.nn.functional as F
import pytorch_lightning as pl

from pricePrediction import config
from pricePrediction.ArgParser_base import ArgParseable
from pricePrediction.config import USE_FEATURES_NET
from pricePrediction.nets.FDS_imbalance import FDS


class PricePredictorModule(pl.LightningModule, ArgParseable):

    DESIRED_PARAMS_TO_ASK= ['n_layers', 'hidden_size_node', 'hidden_size_edges', 'fcc_hidden_size', 'use_fds', 'lr',
                            'dropout', 'gnn_class', 'weight_decay', 'training_loss']
    def __init__(
        self,
        n_layers: int = config.N_LAYERS,
        hidden_size_node: int = config.N_HIDDEN_NODE,
        hidden_size_edges: int = config.N_HIDDEN_EDGE,
        gnn_class="GNN_PNAConv",
        towers: int = 5,
        fcc_hidden_size: int = 25,
        dropout: float = 0,
        use_fds: bool = False,
        training_loss : str= "l2", # l1, l2 or huber
        last_layer_size: int = 1,

        nodes_n_features: int = None,
        edges_n_features: int = None,
        deg: Iterable[float] = None,

        lr: float = config.LEARNING_RATE,
        b1: float = 0.5,
        b2: float = 0.999,
        weight_decay=1e-8,
        data_hparams =None,
        logs_only_in_epoch = False
    ):


        '''

        :param str encodedDir: The directory where the dataset has been prepared as lmdb files. defaults to %(config.ENCODED_DIR)s
        :param str deg_fname: The file containing the statistics about the nodes degree. By default is automatically searched within
                              encodedDir defaults to None
        :param int n_layers: The number of layers.
        :param int hidden_size_node: Channels for nodes.
        :param int hidden_size_edges: Channels for edges.
        :param str gnn_class: Name of the model type.
        :param int towers: Only for gnn_class="GNN_PNAConv", The number of towers of the model.
        :param int fcc_hidden_size: Size of fully connected.
        :param float dropout: Dropout rate.
        :param bool use_fds: Use feature density smoothing for data imbalance
        :param str training_loss: Training loss: "l2" "l1" or "huber"
        :param int last_layer_size: The number of neurons for last layer. Set it to 1 for regression
        :param int nodes_n_features: The number of features used to encode a node
        :param int edges_n_features: The number of features used to encode an edge
        :param Iterable deg: The degree count of the dataset
        :param float lr: Learning rate
        :param float b1: Adam optimizer b1
        :param float b2: Adam optimizer b1
        '''

        super().__init__()
        self.save_hyperparameters() #copies __init__() kwargs to self.hparams
        deg =  self.hparams.deg
        if isinstance(deg, list):
            deg = torch.tensor(deg[:], dtype=torch.long)
        else:
            deg = deg.clone().detach()
        self.hparams.deg = deg



        if gnn_class == "GNN_PNAConv":
            from pricePrediction.nets.basicArchitectures import GNN_PNAConv as GNN
        elif gnn_class == "GNN_AttentiveFP":
            from pricePrediction.nets.basicArchitectures import GNN_AttentiveFP  as GNN
        elif gnn_class == "QdolarAR":
            assert USE_FEATURES_NET
            from pricePrediction.nets.basicArchitectures import QdolarAR as GNN

        else:
            raise ValueError("Error, gnn_class not supported")

        self.net = GNN(** self.hparams)

        self.dropoutLayer = nn.Dropout(self.hparams.dropout)

        self.final_layer = nn.Linear(self.net.fcc_hidden_size, 1)

        if use_fds:
            self.fds = FDS(self.net.latent_size, max_val=10, start_update=1)
        else:
            self.fds = None
        if training_loss=="l1":
            self.trainLossF = F.l1_loss
        elif training_loss=="l2":
            self.trainLossF = F.mse_loss
        elif training_loss=="huber":
            self.trainLossF = F.smooth_l1_loss
        else:
            raise ValueError("Loss option not recognized")

        if logs_only_in_epoch: #TODO: check if this prevents multi-gpu hang
            self.logs_kwargs = dict(on_step=False, on_epoch=True, sync_dist=True)
        else:
            self.logs_kwargs = dict(on_step=True, on_epoch=False, sync_dist=False)


    def compute_training_loss(self, y_pred, y, w=None):
        loss = self.trainLossF(y_pred, y, reduction='none')
        if w is not None:
            loss *= w
        loss = torch.mean(loss)
        return loss

    def forward(self, g):
        return self.compute_y_pred( g, training=False)

    def compute_y_pred(self, g, y=None, training=True):
        x = self.net(g.x, g.edge_index, g.edge_attr, g.batch)
        if training and self.fds is not None:
            self.fds.smooth(x, y, self.current_epoch)
        x = self.dropoutLayer(x)
        y_pred = self.final_layer(x)
        return y_pred.view(-1)

    def resolve_batch(self, batch):
        if isinstance(batch, list):
            return batch[0]
        else:
            return batch

    def training_step(self, batch, batch_idx):
        graphs = self.resolve_batch(batch)
        y_pred = self.compute_y_pred(graphs, training=True)
        loss = self.compute_training_loss(y_pred, graphs.y, graphs.w)

        loss_l1 = F.l1_loss(y_pred, graphs.y)

        self.log('loss', loss, **self.logs_kwargs)
        self.log('loss_l1', loss_l1, prog_bar=True, **self.logs_kwargs)

        tqdm_dict = {'loss': loss.detach()}
        output = OrderedDict({
            'loss': loss,
            'targets': graphs.y,
            'preds': y_pred.detach(),
            'progress_bar': tqdm_dict,
            'log': tqdm_dict
        })
        return output

    def _validation_step(self, graphs, batch_idx):
        with torch.no_grad():
            y_pred = self.compute_y_pred(graphs, training=False)
            loss = self.compute_training_loss(y_pred, graphs.y, graphs.w)
            loss_l1 = F.l1_loss(y_pred, graphs.y)
            return y_pred, loss, loss_l1

    def validation_step(self, batch, batch_idx):
        graphs = self.resolve_batch(batch)
        y_pred, loss, loss_l1 = self._validation_step(graphs, batch_idx)
        # tensorboard = self.logger.experiment
        # tensorboard.add_histogram("val_price", graphs.y)
        # tensorboard.add_histogram("val_preds", y_pred)
        self.log('val_loss', loss, prog_bar=True, sync_dist=True) #rank_zero_only=True
        self.log('val_lossL1', loss_l1, prog_bar=True, sync_dist=True)
        return loss

    def training_epoch_end(self, outputs ):
        # tensorboard = self.logger.experiment
        # tensorboard.add_histogram("train_price", y)
        # tensorboard.add_histogram("train_preds", torch.cat([ x["preds"] for x in  outputs]))
        # tensorboard.add_histogram("train_price", torch.cat([ x["targets"] for x in  outputs]))
        print("Epoch %d done"%self.current_epoch)

    def test_step(self, batch, batch_idx):
        graphs = self.resolve_batch(batch)
        y_pred, loss, loss_l1 = self._validation_step(graphs, batch_idx)
        tensorboard = self.logger.experiment
        # tensorboard.add_histogram("test_price", graphs.y)
        # tensorboard.add_histogram("test_preds", y_pred)
        self.log('test_loss', loss, prog_bar=True, sync_dist=True)
        self.log('test_lossL1', loss_l1, prog_bar=True, sync_dist=True)
        return loss

    def configure_optimizers(self):
        lr = self.hparams.lr
        b1 = self.hparams.b1
        b2 = self.hparams.b2

        opt = torch.optim.Adam(self.parameters(), lr=lr, betas=(b1, b2), weight_decay=self.hparams.weight_decay)

        # for param_group in opt.param_groups:
        #     print(param_group["lr"])

        conf = {
            'optimizer': opt,
            'scheduler': torch.optim.lr_scheduler.ReduceLROnPlateau(opt, verbose=True, cooldown=2,
                                                                patience=config.PATIENT_REDUCE_LR_PLATEAU_N_EPOCHS, ),
            'monitor': 'val_loss'
        }

        return conf

if __name__ == "__main__":
    from torch_geometric.data import Batch

    degree =[0, 41130, 117278, 70152, 3104]

    gnn_class = "QdolarAR" #"GNN_PNAConv"# "GNN_AttentiveFP"

    if gnn_class != "QdolarAR":
        from pricePrediction.preprocessData.smilesToGraph import smiles_to_graph
    else:
        from pricePrediction.preprocessData.smilesToDescriptors import smiles_to_graph

    g = smiles_to_graph("CCCCCCCCCCC")
    nodes_n_features = g.x.shape[1]
    edges_n_features = g.edge_attr.shape[1]

    net = PricePredictorModule(nodes_n_features=nodes_n_features, edges_n_features=edges_n_features,
                               deg=degree, lr=1e-3, gnn_class=gnn_class)

    print( net )

    graphs = [smiles_to_graph(smi) for smi in ["CCCCCO", "CCCCCNCCCCN"]]
    print(graphs[0])
    print( net(Batch.from_data_list(graphs)))