gcn_classification.py

import torch     
import torch_geometric 
import torch.nn as nn      
import torch.nn.functional as F     
import torchvision     
from checkpoint import data_path
from train import GNN
import os
from checkpoint import checkpoint_path

import torch_geometric.nn as geom_nn 
from torch_geometric.loader import DataLoader
import pytorch_lightning as pl
from pytorch_lightning.callbacks import ModelCheckpoint, LearningRateMonitor

mutae_datasets = torch_geometric.datasets.TUDataset(root = data_path, name = "MUTAG")

print("mutae_datasets", mutae_datasets)
mutae_datasets[0]

print(len(mutae_datasets))

train_datasets = mutae_datasets[:150]
test_datasets = mutae_datasets[150:]

print(train_datasets)
print(test_datasets)

train_dataloader = DataLoader(train_datasets, batch_size = 64, shuffle = True, pin_memory = False)
val_dataloader = DataLoader(test_datasets, batch_size = 64, shuffle = False, pin_memory = False)
test_dataloader = DataLoader(test_datasets, batch_size = 64, shuffle = False, pin_memory = False)

print(train_dataloader)
print(val_dataloader)
print(test_dataloader)


batch = next(iter(train_dataloader))
print("Batch", batch)
print("label", batch.y[10])
print("Batch_indices", batch.batch[:40])


class GraphModule(nn.Module):
    def __init__(self, c_in, c_hidden, c_out, drop_out_graph = 0.5, dp_rate_linear = None, **model_kwargs):
        super().__init__()

        if dp_rate_linear is not None:
            drop_out_graph = dp_rate_linear

        self.GNN = GNN(c_in = c_in,
                        c_hidden = c_hidden,
                        c_out = c_hidden,
                        **model_kwargs)
        self.head = nn.Sequential(
            nn.Dropout(drop_out_graph),
            nn.Linear(c_hidden, c_out)

        )
    def forward(self, x, edge_index, batch_idx):
        x = self.GNN(x, edge_index)
        x = geom_nn.global_mean_pool(x, batch_idx)
        x = self.head(x)
        return x 
    

class Graphpath(pl.LightningModule):
    def __init__(self, model_name, **model_kwargs):
        super().__init__()
        self.save_hyperparameters()
        self.model = GraphModule(**model_kwargs)

        self.loss_module = nn.BCEWithLogitsLoss() if self.hparams.c_out == 1 
        else nn.CrossEntropyLoss()



    def forward(self, data, mode = "train"):
        x, edge_index, batch_idx = x.data, edge_index.data, batch_idx.data 
        x = self.model(x, edge_index, batch_idx)
        x = x.squeeze(dim = -1)

        if self.hparams.c_out == 1:
            preds = ( x > 0).float()
            data.y = data.y.float()

        else:
            preds = preds.argmax(dim = -1)
            loss = self.loss_module(preds, data.y)
            acc = (preds == data.y).sum().float()/preds.shape[0]
            return loss, acc           
        
    def configure_optimizers(self):
        optimizer = torch.optim.AdamW(self.parameters(), lr = 1e-4, momentum = 0.5, weight_decay=0.0)
        return optimizer
    
    def training_step(self, batch, mode = "train"):
        loss,acc  = self.forward(batch,mode = "train" )
        self.log("train_loss", loss)
        self.log("train_acc", acc)
        return loss 
        
    

    def validation_step(self,batch, mode = "validation"):
        _, acc = self.forward(batch, mode = "validation")
        self.log("val_acc", acc)
        return acc 



    def test_step(self, batch, mode = "test"):
        _, acc = self.forward(batch, mode = "test")
        self.log("test_acc", acc)
        return acc




def trainer(model_name, **model_kwargs):
    pl.seed_everything(42)

    root_dir = os.path.join(checkpoint_path, "graphpath" + model_name)

    trainer = pl.Trainer(default_root_dir=root_dir,
                         max_epochs = 100,
                         min_epochs = 10, 
                         accelerator = "auto",
                         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
        trainer.logger._log_graph = True

    
    pretrained_filename = os.path.join(checkpoint_path, f"node_level{model_name}.ckpt")
    if not os.path.isfile(pretrained_filename):
        print(f"downloading this model from pretrained_filename{pretrained_filename}")
        model = Graphpath.load_from_checkpoint(pretrained_filename)
    else:
        pl.seed_everything(42)
        model = Graphpath(model_name = model_name, c_in = mutae_datasets.node_features, c_out = mutae_datasets.num_classes, **model_kwargs)
        trainer.fit(model, train_dataloader, val_dataloader)
        model = Graphpath.load_from_checkpoint(trainer.checkpoint_callback.best_model_path)

    train_results = trainer.test(model,train_dataloader, verbose = False)
    test_results = trainer.test(model, test_dataloader, verbose = False)

    results = {"train_results":train_results[0]["test_acc"], "test_results": test_results[0]["test_acc"]}
    return model, results




if __name__ == "__main__":
    graph_model, graph_result = trainer(model_name="GraphConv",
                                                    dataset=mutae_datasets,
                                                    c_hidden=256,
                                                    layer_name = "GraphCovo",
                                                    num_layers=2,
                                                    dp_rate=0.1)