glam4cm 0.1.0__py3-none-any.whl

glam4cm/trainers/gnn_edge_classifier.py

@@ -0,0 +1,126 @@
+from typing import List
+import torch
+from collections import defaultdict
+from torch_geometric.loader import DataLoader
+from glam4cm.data_loading.data import GraphData
+from glam4cm.models.gnn_layers import (
+    GNNConv, 
+    EdgeClassifer
+)
+
+from glam4cm.trainers.gnn_trainer import Trainer
+from glam4cm.settings import device
+
+
+class GNNEdgeClassificationTrainer(Trainer):
+    """
+    Trainer class for GNN Link Prediction
+    This class is used to train the GNN model for the link prediction task
+    The model is trained to predict the link between two nodes
+    """
+    def __init__(
+            self, 
+            model: GNNConv, 
+            predictor: EdgeClassifer, 
+            dataset: List[GraphData],
+            cls_label='type',
+            lr=1e-3,
+            num_epochs=100,
+            batch_size=32,
+            use_edge_attrs=False,
+            logs_dir='./logs'
+        ) -> None:
+
+        super().__init__(
+            model=model,
+            predictor=predictor,
+            cls_label=cls_label,
+            lr=lr,
+            num_epochs=num_epochs,
+            use_edge_attrs=use_edge_attrs,
+            logs_dir=logs_dir
+        )
+
+        self.dataloader = DataLoader(
+            dataset, 
+            batch_size=batch_size, 
+            shuffle=True
+        )
+
+        print("GNN Trainer initialized.")
+
+
+    def train(self):
+        self.model.train()
+        self.predictor.train()
+
+        all_preds, all_labels = list(), list()
+        epoch_loss = 0
+        epoch_metrics = defaultdict(float)
+        for data in self.dataloader:
+            self.optimizer.zero_grad()
+            self.model.zero_grad()
+            self.predictor.zero_grad()
+            x = data.x
+            edge_index =  data.train_pos_edge_label_index
+            train_mask = data.train_edge_mask
+            edge_attr = data.edge_attr[train_mask] if self.use_edge_attrs else None
+            
+            h = self.get_logits(x, edge_index, edge_attr)
+            scores = self.get_prediction_score(h, edge_index, edge_attr)
+            labels = getattr(data, f"edge_{self.cls_label}")[train_mask]
+            loss = self.criterion(scores, labels.to(device))
+            all_preds.append(scores.detach().cpu())
+            all_labels.append(labels)
+
+            loss.backward()
+            self.optimizer.step()
+            self.scheduler.step()
+            epoch_loss += loss.item()
+                        
+        
+        all_preds = torch.cat(all_preds, dim=0)
+        all_labels = torch.cat(all_labels, dim=0)
+        epoch_metrics = self.compute_metrics(all_preds, all_labels)
+        epoch_metrics['loss'] = epoch_loss        
+        epoch_metrics['phase'] = 'train'
+
+        # print(f"Train Metrics: {epoch_metrics}")
+
+        return epoch_metrics
+
+
+    def test(self):
+        self.model.eval()
+        self.predictor.eval()
+        all_preds, all_labels = list(), list()
+        with torch.no_grad():
+            epoch_loss = 0
+            epoch_metrics = defaultdict(float)
+            for data in self.dataloader:
+                x = data.x
+                edge_index =  data.test_pos_edge_label_index
+                test_mask = data.test_edge_mask
+                edge_attr = data.edge_attr[test_mask] if self.use_edge_attrs else None
+                
+                h = self.get_logits(x, edge_index, edge_attr)
+                scores = self.get_prediction_score(h, edge_index, edge_attr)
+                labels = getattr(data, f"edge_{self.cls_label}")[test_mask]
+                all_preds.append(scores.detach().cpu())
+                all_labels.append(labels)
+                loss = self.criterion(scores, labels.to(device))
+                
+                epoch_loss += loss.item()
+
+            all_preds = torch.cat(all_preds, dim=0)
+            all_labels = torch.cat(all_labels, dim=0)
+            epoch_metrics = self.compute_metrics(all_preds, all_labels)
+            
+            epoch_metrics['loss'] = epoch_loss
+            epoch_metrics['phase'] = 'test'
+            # print(f"Epoch Test Loss: {epoch_loss}\nTest Accuracy: {epoch_acc}\nTest F1: {epoch_f1}")
+            self.results.append(epoch_metrics)
+
+            print(f"Epoch: {len(self.results)}\n{epoch_metrics}")
+        
+        return epoch_metrics

glam4cm/trainers/gnn_graph_classifier.py

@@ -0,0 +1,123 @@
+from typing import List, Tuple
+import torch
+from collections import defaultdict
+from torch_geometric.loader import DataLoader
+
+from torch_geometric.data import Data
+from glam4cm.models.gnn_layers import (
+    GNNConv,
+    GraphClassifer
+)
+from glam4cm.trainers.gnn_trainer import Trainer
+from glam4cm.settings import device
+
+
+class GNNGraphClassificationTrainer(Trainer):
+    """
+    Trainer class for GNN Graph Classfication
+    This class is used to train the GNN model for the link prediction task
+    The model is trained to predict the link between two nodes
+    """
+    def __init__(
+            self, 
+            model: GNNConv, 
+            predictor: GraphClassifer,
+            dataset: List[Tuple[Data, Data]],
+            cls_label='label',
+            lr=1e-4,
+            num_epochs=100,
+            batch_size=32,
+            use_edge_attrs=False,
+            logs_dir='./logs'
+        ) -> None:
+
+        super().__init__(
+            model=model,
+            predictor=predictor,
+            cls_label='type',
+            lr=lr,
+            num_epochs=num_epochs,
+            use_edge_attrs=use_edge_attrs,
+            logs_dir=logs_dir
+        )
+
+        self.cls_label = cls_label
+        self.dataloaders = dict()
+        self.dataloaders['train'] = DataLoader(dataset['train'], batch_size=batch_size, shuffle=True)
+        self.dataloaders['test'] = DataLoader(dataset['test'], batch_size=batch_size, shuffle=False)
+
+        self.results = list()
+
+        print("GNN Trainer initialized.")
+
+
+    def train(self):
+        self.model.train()
+        self.predictor.train()
+
+        epoch_loss = 0
+        epoch_metrics = defaultdict(float)
+        preds, all_labels = list(), list()
+        for data in self.dataloaders['train']:
+            self.optimizer.zero_grad()
+            self.model.train()
+            self.predictor.train()
+            
+            h = self.model(data.x.to(device), data.edge_index.to(device))
+            g_pred = self.predictor(h, data.batch.to(device))
+
+            
+            labels = getattr(data, f"graph_{self.cls_label}")
+            loss = self.criterion(g_pred, labels.to(device))
+
+            preds.append(g_pred.detach().cpu())
+            all_labels.append(labels)
+
+            loss.backward()
+            self.optimizer.step()
+            self.scheduler.step()
+            epoch_loss += loss.item()
+
+        
+        preds = torch.cat(preds, dim=0)
+        labels = torch.cat(all_labels, dim=0)
+        
+        epoch_metrics = self.compute_metrics(preds, labels)
+        epoch_metrics['loss'] = epoch_loss
+        epoch_metrics['phase'] = 'train'
+
+        self.results.append(epoch_metrics)
+
+        return epoch_metrics
+
+
+    def test(self):
+        self.model.eval()
+        self.predictor.eval()
+        with torch.no_grad():
+            epoch_loss = 0
+            preds, all_labels = list(), list()
+            for data in self.dataloaders['test']:
+                h = self.model(data.x.to(device), data.edge_index.to(device))
+                g_pred = self.predictor(h, data.batch.to(device))
+                labels = getattr(data, f"graph_{self.cls_label}")
+
+                loss = self.criterion(g_pred, labels.to(device))
+                epoch_loss += loss.item()
+
+                preds.append(g_pred.cpu().detach())
+                all_labels.append(labels.cpu())
+
+            
+            preds = torch.cat(preds, dim=0)
+            labels = torch.cat(all_labels, dim=0)
+
+            epoch_metrics = self.compute_metrics(preds, labels)
+            epoch_metrics['loss'] = epoch_loss
+            epoch_metrics['phase'] = 'test'
+            self.results.append(epoch_metrics)
+        
+        s2t = lambda x: x.replace("_", " ").title()
+        print(f"Epoch: {len(self.results)//2} {' | '.join([f'{s2t(k)}: {v:.4f}' for k, v in epoch_metrics.items() if k != 'phase'])}")
+
+        return epoch_metrics

glam4cm/trainers/gnn_link_predictor.py

@@ -0,0 +1,144 @@
+from torch_geometric.loader import DataLoader
+from torch_geometric.data import Data
+import torch
+from collections import defaultdict
+from typing import List
+
+from glam4cm.models.gnn_layers import (
+    GNNConv, 
+    EdgeClassifer
+)
+
+from glam4cm.trainers.gnn_trainer import Trainer
+from tqdm.auto import tqdm
+from glam4cm.settings import device
+
+
+class GNNLinkPredictionTrainer(Trainer):
+    """
+    Trainer class for GNN Link Prediction
+    This class is used to train the GNN model for the link prediction task
+    The model is trained to predict the link between two nodes
+    """
+    def __init__(
+            self, 
+            model: GNNConv, 
+            predictor: EdgeClassifer, 
+            dataset: List[Data],
+            cls_label='type',
+            lr=1e-3,
+            num_epochs=100,
+            batch_size=32,
+            use_edge_attrs=False,
+            logs_dir='./logs'
+        ) -> None:
+
+        super().__init__(
+            model=model,
+            predictor=predictor,
+            lr=lr,
+            cls_label=cls_label,
+            num_epochs=num_epochs,
+            use_edge_attrs=use_edge_attrs,
+            logs_dir=logs_dir
+        )        
+        self.dataloader = DataLoader(
+            dataset, batch_size=batch_size, shuffle=True
+        )
+        self.results = list()
+
+        print("GNN Trainer initialized.")
+
+
+
+    def train(self):
+        self.model.train()
+        self.predictor.train()
+
+        all_preds, all_labels = list(), list()
+        epoch_loss = 0
+        epoch_metrics = defaultdict(float)
+        for data in tqdm(self.dataloader, desc='Training Batches'):
+            self.optimizer.zero_grad()
+            self.model.zero_grad()
+            self.predictor.zero_grad()
+
+            x = data.x
+            pos_edge_index =  data.train_pos_edge_label_index
+            neg_edge_index = data.train_neg_edge_label_index
+            train_mask = data.train_edge_mask
+            edge_attr = data.edge_attr[train_mask] if self.use_edge_attrs else None
+            
+            h = self.get_logits(x, pos_edge_index, edge_attr)
+            # h = x
+
+            pos_scores = self.get_prediction_score(h, pos_edge_index, edge_attr)
+            neg_scores = self.get_prediction_score(h, neg_edge_index, edge_attr)
+            loss = self.compute_loss(pos_scores, neg_scores)
+            all_labels.append(torch.cat([torch.ones(pos_scores.size(0)), torch.zeros(neg_scores.size(0))]))
+            all_preds.append(torch.cat([pos_scores.detach().cpu(), neg_scores.detach().cpu()]))
+
+            loss.backward()
+            self.optimizer.step()
+            self.scheduler.step()
+            epoch_loss += loss.item()
+                        
+        
+        all_preds = torch.cat(all_preds, dim=0)
+        all_labels = torch.cat(all_labels, dim=0)
+        epoch_metrics = self.compute_metrics(all_preds, all_labels)
+        epoch_metrics['loss'] = epoch_loss        
+        epoch_metrics['phase'] = 'train'
+
+        return epoch_metrics
+
+
+    def test(self):
+        self.model.eval()
+        self.predictor.eval()
+        all_preds, all_labels = list(), list()
+        with torch.no_grad():
+            epoch_loss = 0
+            epoch_metrics = defaultdict(float)
+            for data in tqdm(self.dataloader, desc='Testing Batches'):
+                
+                x = data.x
+                pos_edge_index =  data.test_pos_edge_label_index
+                neg_edge_index = data.test_neg_edge_label_index
+                test_mask = data.test_edge_mask
+                edge_attr = data.edge_attr[test_mask] if self.use_edge_attrs else None
+
+
+                h = self.get_logits(x, pos_edge_index, edge_attr)
+                # h = x
+                pos_score = self.get_prediction_score(h, pos_edge_index, edge_attr)
+                neg_score = self.get_prediction_score(h, neg_edge_index, edge_attr)
+
+                loss = self.compute_loss(pos_score, neg_score)
+                all_labels.append(torch.cat([torch.ones(pos_score.size(0)), torch.zeros(neg_score.size(0))]))
+                all_preds.append(torch.cat([pos_score.detach().cpu(), neg_score.detach().cpu()]))
+
+                epoch_loss += loss.item()
+
+            all_preds = torch.cat(all_preds, dim=0)
+            all_labels = torch.cat(all_labels, dim=0)
+            epoch_metrics = self.compute_metrics(all_preds, all_labels)
+
+            epoch_metrics['loss'] = epoch_loss
+            epoch_metrics['phase'] = 'test'
+            # print(f"Epoch Test Loss: {epoch_loss}\nTest Accuracy: {epoch_acc}\nTest F1: {epoch_f1}")
+            self.results.append(epoch_metrics)
+
+            print(f"Test Epoch: {len(self.results)}\n{epoch_metrics}")
+
+        return epoch_metrics        
+
+    def compute_loss(self, pos_score, neg_score):
+        pos_label = torch.ones(pos_score.size(0), dtype=torch.long).to(device)
+        neg_label = torch.zeros(neg_score.size(0), dtype=torch.long).to(device)
+
+        scores = torch.cat([pos_score, neg_score], dim=0)
+        labels = torch.cat([pos_label, neg_label], dim=0)
+
+        loss = self.criterion(scores, labels)
+        return loss

glam4cm/trainers/gnn_node_classifier.py

@@ -0,0 +1,135 @@
+from typing import List
+from torch_geometric.loader import DataLoader
+import torch
+from collections import defaultdict
+from torch_geometric.data import Data
+from glam4cm.models.gnn_layers import (
+    GNNConv, 
+    NodeClassifier
+)
+from glam4cm.trainers.gnn_trainer import Trainer
+from glam4cm.settings import device
+
+
+class GNNNodeClassificationTrainer(Trainer):
+    """
+    Trainer class for GNN Link Prediction
+    This class is used to train the GNN model for the link prediction task
+    The model is trained to predict the link between two nodes
+    """
+    def __init__(
+            self, 
+            model: GNNConv, 
+            predictor: NodeClassifier, 
+            dataset: List[Data],
+            cls_label,
+            exclude_labels=None,
+            lr=1e-3,
+            num_epochs=100,
+            batch_size=32,
+            use_edge_attrs=False,
+            logs_dir='./logs'
+        ) -> None:
+
+        super().__init__(
+            model=model,
+            predictor=predictor,
+            cls_label=cls_label,
+            lr=lr,
+            num_epochs=num_epochs,
+            use_edge_attrs=use_edge_attrs,
+            logs_dir=logs_dir
+        )
+
+        self.exclude_labels = torch.tensor(exclude_labels, dtype=torch.long)
+        self.results = list()
+        self.dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
+        print("GNN Trainer initialized.")
+
+
+
+    def train(self):
+        self.model.train()
+        self.predictor.train()
+
+        all_preds, all_labels = list(), list()
+        epoch_loss = 0
+        epoch_metrics = defaultdict(float)
+        # for i, data in tqdm(enumerate(self.dataloader), desc=f"Training batches", total=len(self.dataloader)):
+        for data in self.dataloader:
+            self.optimizer.zero_grad()
+            self.model.zero_grad()
+            self.predictor.zero_grad()
+            
+            h = self.get_logits(
+                data.x, 
+                data.edge_index,
+                data.edge_attr if self.use_edge_attrs else None
+            )
+            scores = self.get_prediction_score(h)[data.train_node_mask]
+            labels = getattr(data, f"node_{self.cls_label}")[data.train_node_mask]
+            
+            mask = ~torch.isin(labels, self.exclude_labels)
+            labels = labels[mask]
+            scores = scores[mask]
+
+            loss = self.criterion(scores, labels.to(device))
+            
+            all_preds.append(scores.detach().cpu())
+            all_labels.append(labels)
+
+            loss.backward()
+            self.optimizer.step()
+            self.scheduler.step()
+            epoch_loss += loss.item()
+                        
+        
+        all_preds = torch.cat(all_preds, dim=0)
+        all_labels = torch.cat(all_labels, dim=0)
+        epoch_metrics = self.compute_metrics(all_preds, all_labels)
+        epoch_metrics['loss'] = epoch_loss        
+        epoch_metrics['phase'] = 'train'
+
+        return epoch_metrics
+
+
+    def test(self):
+        self.model.eval()
+        self.predictor.eval()
+        all_preds, all_labels = list(), list()
+        with torch.no_grad():
+            epoch_loss = 0
+            epoch_metrics = defaultdict(float)
+            # for _, data in tqdm(enumerate(self.dataloader), desc=f"Evaluating batches", total=len(self.dataloader)):
+            for data in self.dataloader:
+                h = self.get_logits(
+                    data.x, 
+                    data.edge_index,
+                    data.edge_attr if self.use_edge_attrs else None
+                )
+                
+                scores = self.get_prediction_score(h)[data.test_node_mask]
+                labels = getattr(data, f"node_{self.cls_label}")[data.test_node_mask]
+
+                mask = ~torch.isin(labels, self.exclude_labels)
+                labels = labels[mask]
+                scores = scores[mask]
+                loss = self.criterion(scores, labels.to(device))
+                epoch_loss += loss.item()
+
+
+                all_preds.append(scores.detach().cpu())
+                all_labels.append(labels)
+                
+
+            all_preds = torch.cat(all_preds, dim=0)
+            all_labels = torch.cat(all_labels, dim=0)
+            epoch_metrics = self.compute_metrics(all_preds, all_labels)
+            
+            epoch_metrics['loss'] = epoch_loss
+            epoch_metrics['phase'] = 'test'
+            self.results.append(epoch_metrics)
+
+            print(f"Epoch: {len(self.results)}\n{epoch_metrics}")
+
+        return epoch_metrics

glam4cm/trainers/gnn_trainer.py

@@ -0,0 +1,129 @@
+from abc import abstractmethod
+import torch
+from typing import Union
+import pandas as pd
+
+from glam4cm.models.gnn_layers import (
+    GNNConv, 
+    EdgeClassifer,
+    NodeClassifier
+)
+from glam4cm.settings import device
+from itertools import chain
+from tqdm.auto import tqdm
+import torch.nn as nn
+from torch.optim.lr_scheduler import CosineAnnealingLR
+from torch.optim import Adam
+from torch_geometric.loader import DataLoader
+
+from tensorboardX import SummaryWriter
+from glam4cm.trainers.metrics import compute_classification_metrics
+
+
+class Trainer:
+    """
+    Trainer class for GNN Link Prediction
+    This class is used to train the GNN model for the link prediction task
+    The model is trained to predict the link between two nodes
+    """
+    def __init__(
+            self, 
+            model: GNNConv, 
+            predictor: Union[EdgeClassifer, NodeClassifier], 
+            cls_label,
+            lr=1e-3,
+            num_epochs=100,
+            use_edge_attrs=False,
+
+            logs_dir='./logs'
+        ) -> None:
+        self.model = model
+        self.predictor = predictor
+        self.model.to(device)
+        self.predictor.to(device)
+
+        self.cls_label = cls_label
+        self.num_epochs = num_epochs
+        
+        self.optimizer = Adam(chain(model.parameters(), predictor.parameters()), lr=lr)
+        self.scheduler = CosineAnnealingLR(self.optimizer, T_max=num_epochs)
+        
+        self.edge2index = lambda g: torch.stack(list(g.edges())).contiguous()
+        self.results = list()
+        self.criterion = nn.CrossEntropyLoss()
+
+        self.use_edge_attrs = use_edge_attrs
+
+        self.logs_dir = logs_dir
+
+        self.writer = SummaryWriter(log_dir=self.logs_dir)
+
+        print("GNN Trainer initialized.")
+
+
+    def set_dataloader(self, dataset, batch_size):
+        self.dataloader = DataLoader(
+            dataset, batch_size=batch_size, shuffle=True
+        )
+
+    @abstractmethod
+    def train(self):
+        pass
+
+    @abstractmethod
+    def test(self):
+        pass
+
+
+    def get_logits(self, x, edge_index, edge_attr=None):
+        edge_index = edge_index.to(device)
+        x = x.to(device)
+
+        if edge_attr is not None:
+            edge_attr = edge_attr.to(device)
+            h = self.model(x, edge_index, edge_attr)
+        else:
+            h = self.model(x, edge_index)
+        return h
+    
+
+    def get_prediction_score(self, h, edge_index=None, edge_attr=None):
+        h = h.to(device)
+        if edge_attr is not None:
+            edge_attr = edge_attr.to(device)
+            edge_index = edge_index.to(device)
+            prediction_score = self.predictor(h, edge_index, edge_attr)
+        elif edge_index is not None:
+            edge_index = edge_index.to(device)
+            prediction_score = self.predictor(h, edge_index)
+        else:
+            prediction_score = self.predictor(h)
+        return prediction_score
+        
+    
+    def plot_metrics(self):
+        results = pd.DataFrame(self.results)
+        df = pd.DataFrame(results, index=range(1, len(results)+1))
+        df['epoch'] = df.index
+
+        columns = [c for c in df.columns if c not in ['epoch', 'phase']]
+        df.loc[df['phase'] == 'test'].plot(x='epoch', y=columns, kind='line')
+
+
+    def run(self):
+        for epoch in tqdm(range(self.num_epochs), desc="Running Epochs"):
+            train_metrics = self.train()
+            test_metrics = self.test()
+
+            for k, v in train_metrics.items():
+                if k != 'phase':
+                    self.writer.add_scalar(f"train/{k}", v, epoch)
+            
+            for k, v in test_metrics.items():
+                if k != 'phase':
+                    self.writer.add_scalar(f"test/{k}", v, epoch)
+    
+        self.writer.close()
+    
+    def compute_metrics(self, all_preds, all_labels):
+        return compute_classification_metrics(all_preds, all_labels)