train.py

import torch   
import torch_geometric
import torch_geometric.nn as geom_nn
import torch_geometric.data as geom_data
from checkpoint import data_path
import torch.nn as nn 
import pytorch_lightning as pl 
from torch_geometric.loader import DataLoader
import os
from checkpoint import checkpoint_path
from pytorch_lightning.callbacks import ModelCheckpoint, LearningRateMonitor
gnn_geometric = {
    "GCN":geom_nn.GCNConv,
    "GAT":geom_nn.GATConv,
    "Graph":geom_nn.GraphConv,
    "GraphConv":geom_nn.GraphConv
}

gnn_geometric

cora_datasets = torch_geometric.datasets.Planetoid(root = data_path,  name = "cora")

cora_datasets[0]


class GNN(nn.Module):
    def __init__(self, c_in, c_hidden, c_out, num_layers = 2, layer_name = "GCN", dp_out = 0.2, dp_rate = None, **kwargs):
        super().__init__()

        if dp_rate is not None:
            dp_out = dp_rate

        layers = []

        gnn_layers = gnn_geometric[layer_name]
        in_channels = c_in
        out_channels = c_hidden

        for idx in range(num_layers -1):

            layers += [
                gnn_layers(in_channels = in_channels,
                           out_channels = out_channels,
                           **kwargs),
                           nn.ReLU(inplace = True),
                           nn.Dropout(dp_out),

            ]


            in_channels = c_hidden

        layers += [     
            gnn_layers (in_channels = in_channels,
                        out_channels = c_out,
                        **kwargs)
        ]

        self.layers = nn.ModuleList(layers)

    def forward(self, x, edgelist):
        for layers in self.layers:
            if isinstance(layers, geom_nn.MessagePassing):

                x = layers(x, edgelist)
            else:
                x = layers(x)
        return x 
    




class MLP(nn.Module):
    def __init__(self, c_in, c_hidden, c_out, num_layers = 2, dp_out = 0.2, dp_rate = None, layer_name = "GCN", **kwargs):
        super().__init__()

        if dp_rate is not None:
            dp_out = dp_rate

        layers = []

        in_channels = c_in 
        out_channels = c_hidden 

        for idx in range(num_layers-1):

            layers += [nn.Linear(in_channels, out_channels,**kwargs),
                       nn.ReLU(inplace = True),
                       nn.Dropout(dp_out),
                       ]
            

            in_channels = c_hidden
            out_channels = c_out

        layers += [
            nn.Linear(in_channels,
                      out_channels,
                      ),

        ]

        self.layers = nn.Sequential(*layers)



    def forward(self, x, *args, **kwargs):

        return self.layers(x)
    




class GCN(pl.LightningModule):
    def __init__(self, model_name, **model_kwargs):
        super().__init__()
        self.save_hyperparameters()

        if model_name == "MLP":
            self.model = MLP(**model_kwargs)
        elif model_name == "GNN":
            self.model = GNN(**model_kwargs)
        else:
            raise ValueError(f"Unknown model_name: {model_name}")

        self.loss_module = nn.CrossEntropyLoss()



    def forward(self, data, mode = "train"):
        x, edge_index = data.x,  data.edge_index 

        x = self.model(x, edge_index)

        if mode == "train":
            mask = data.train_mask

        elif mode == "val":
            mask = data.val_mask
        elif mode == "test":
            mask = data.test_mask

        else:
            raise ValueError(f"unknown mode of the train value {mode}")
        loss = self.loss_module(x[mask], data.y[mask])
        acc = (x[mask].argmax(dim = -1)==data.y[mask]).float().mean()
        return loss, acc 
    
    def configure_optimizers(self):
        optimizer = torch.optim.SGD(self.parameters(), lr = 0.1, momentum = 0.9, weight_decay = 2e-3)
        return optimizer
    
    def training_step(self, batch, batch_idx):
        loss, acc = self.forward(batch, mode = "train")
        self.log("train_loss", loss, batch_size = 1)
        self.log("train_acc", acc, batch_size = 1)
        return loss
    
    def validation_step(self, batch, batch_idx):
       _, acc = self.forward(batch, mode = "val")
       self.log("val_acc", acc, batch_size = 1)
       return acc            
    
    def test_step(self, batch, batch_idx):
        _,acc = self.forward(batch, mode = "test")
        self.log("test_acc", acc, batch_size = 1)
        return acc
    



def training_step(model_name, dataset, **model_kwargs):
    pl.seed_everything(42)
    geom_data_loader = DataLoader(dataset, batch_size = 1)

    root_dir = os.path.join(checkpoint_path, "graph_data" + model_name)
    os.makedirs(root_dir, exist_ok = True)

    trainer = pl.Trainer(default_root_dir = root_dir,
                         accelerator = "auto",
                         max_epochs = 100,
                         min_epochs = 30,
                         callbacks =[ModelCheckpoint(save_weights_only = True, mode = "max", monitor = "val_acc"), LearningRateMonitor("epoch")])
    
    if trainer.logger is not None:
        trainer.logger._default_hp_metrics = None



    pretrained_filename = os.path.join(checkpoint_path, f"NodeLevel{model_name}.ckpt")
    if os.path.isfile(pretrained_filename):
        print(f"loading pretrained checkpoint from {pretrained_filename}")
        model = GCN.load_from_checkpoint(pretrained_filename)
    else:
        pl.seed_everything(42)
        model = GCN(model_name = model_name, c_in = dataset.num_node_features, c_out = dataset.num_classes, **model_kwargs)
        trainer.fit(model, geom_data_loader, geom_data_loader)
        model = GCN.load_from_checkpoint(trainer.checkpoint_callback.best_model_path)


    test_result = trainer.test(model, geom_data_loader, verbose = False)
    batch = next(iter(geom_data_loader))
    _, train_acc = model.forward(batch, mode = "train")
    _, val_acc = model.forward(batch, mode = "val")

    results = {"train_acc": train_acc,
               "val_acc": val_acc,
               "test_acc": test_result[0]["test_acc"]}
    
    return model, results
def print_results(result_dict):
    if "train_acc" in result_dict:
        print(f"Train accuracy: {(100.0*result_dict['train_acc']):4.2f}%")
    if "val_acc" in result_dict:
        print(f"Val accuracy:   {(100.0*result_dict['val_acc']):4.2f}%")
    print(f"Test accuracy:  {(100.0*result_dict['test_acc']):4.2f}%")




if __name__ == "__main__":
    node_mlp_model, node_mlp_result = training_step(model_name="MLP",
                                                    dataset=cora_datasets,
                                                    c_hidden=16,
                                                    num_layers=2,
                                                    dp_rate=0.1)

    print_results(node_mlp_result)

    node_gnn_model, node_gnn_result = training_step(model_name="GNN",
                                                    layer_name="GCN",
                                                    dataset=cora_datasets,
                                                    c_hidden=16,
                                                    num_layers=2,
                                                    dp_rate=0.1)
    print_results(node_gnn_result)