gaitsetpy.classification.models.gnn

  1import torch
  2import torch.nn as nn
  3import torch.optim as optim
  4import numpy as np
  5from typing import List, Dict, Any, Optional, Union
  6from ...core.base_classes import BaseClassificationModel
  7from ..utils.preprocess import preprocess_features
  8from sklearn.model_selection import train_test_split
  9from sklearn.metrics import accuracy_score, confusion_matrix, classification_report
 10
 11class SimpleGCN(nn.Module):
 12    def __init__(self, input_dim, hidden_dim, output_dim):
 13        super(SimpleGCN, self).__init__()
 14        self.fc1 = nn.Linear(input_dim, hidden_dim)
 15        self.fc2 = nn.Linear(hidden_dim, output_dim)
 16    def forward(self, x, adj):
 17        h = torch.relu(self.fc1(torch.matmul(adj, x)))
 18        out = self.fc2(torch.matmul(adj, h))
 19        return out
 20
 21class GNNModel(BaseClassificationModel):
 22    """
 23    Simple Graph Neural Network (GCN) classification model using PyTorch.
 24    Implements the BaseClassificationModel interface.
 25    Expects features as node features and adjacency matrix in kwargs.
 26    """
 27    def __init__(self, input_dim=10, hidden_dim=32, output_dim=2, lr=0.001, epochs=20, device=None):
 28        super().__init__(
 29            name="gnn",
 30            description="Graph Convolutional Network (GCN) classifier for gait data classification"
 31        )
 32        self.config = {
 33            'input_dim': input_dim,
 34            'hidden_dim': hidden_dim,
 35            'output_dim': output_dim,
 36            'lr': lr,
 37            'epochs': epochs
 38        }
 39        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
 40        self.model = SimpleGCN(input_dim, hidden_dim, output_dim).to(self.device)
 41        self.epochs = epochs
 42        self.trained = False
 43        self.feature_names = []
 44        self.class_names = []
 45
 46    def train(self, features: List[Dict], **kwargs):
 47        X, y = preprocess_features(features)
 48        # X: (num_nodes, num_features), y: (num_nodes,)
 49        adj = kwargs.get('adjacency_matrix')
 50        if adj is None:
 51            raise ValueError("Adjacency matrix must be provided as 'adjacency_matrix' in kwargs for GNN training.")
 52        X = torch.tensor(X, dtype=torch.float32).to(self.device)
 53        y = torch.tensor(y, dtype=torch.long).to(self.device)
 54        adj = torch.tensor(adj, dtype=torch.float32).to(self.device)
 55        self.feature_names = [f"feature_{i}" for i in range(X.shape[1])]
 56        self.class_names = list(set(y.cpu().numpy()))
 57        criterion = nn.CrossEntropyLoss()
 58        optimizer = optim.Adam(self.model.parameters(), lr=self.config['lr'])
 59        for epoch in range(self.epochs):
 60            self.model.train()
 61            optimizer.zero_grad()
 62            outputs = self.model(X, adj)
 63            loss = criterion(outputs, y)
 64            loss.backward()
 65            optimizer.step()
 66            if (epoch+1) % 5 == 0 or epoch == 0:
 67                print(f"Epoch [{epoch+1}/{self.epochs}], Loss: {loss.item():.4f}")
 68        self.trained = True
 69        print("GNN model trained successfully.")
 70
 71    def predict(self, features: List[Dict], **kwargs) -> np.ndarray:
 72        if not self.trained:
 73            raise ValueError("Model must be trained before making predictions")
 74        X, _ = preprocess_features(features)
 75        adj = kwargs.get('adjacency_matrix')
 76        if adj is None:
 77            raise ValueError("Adjacency matrix must be provided as 'adjacency_matrix' in kwargs for GNN prediction.")
 78        X = torch.tensor(X, dtype=torch.float32).to(self.device)
 79        adj = torch.tensor(adj, dtype=torch.float32).to(self.device)
 80        self.model.eval()
 81        with torch.no_grad():
 82            outputs = self.model(X, adj)
 83            _, predicted = torch.max(outputs.data, 1)
 84        return predicted.cpu().numpy()
 85
 86    def evaluate(self, features: List[Dict], **kwargs) -> Dict[str, float]:
 87        if not self.trained:
 88            raise ValueError("Model must be trained before evaluation")
 89        X, y = preprocess_features(features)
 90        adj = kwargs.get('adjacency_matrix')
 91        if adj is None:
 92            raise ValueError("Adjacency matrix must be provided as 'adjacency_matrix' in kwargs for GNN evaluation.")
 93        X = torch.tensor(X, dtype=torch.float32).to(self.device)
 94        y = np.array(y)
 95        adj = torch.tensor(adj, dtype=torch.float32).to(self.device)
 96        self.model.eval()
 97        with torch.no_grad():
 98            outputs = self.model(X, adj)
 99            _, y_pred = torch.max(outputs.data, 1)
100        y_pred = y_pred.cpu().numpy()
101        accuracy = accuracy_score(y, y_pred)
102        conf_matrix = confusion_matrix(y, y_pred)
103        metrics = {
104            'accuracy': accuracy,
105            'confusion_matrix': conf_matrix.tolist()
106        }
107        detailed_report = kwargs.get('detailed_report', False)
108        if detailed_report:
109            class_report = classification_report(y, y_pred, output_dict=True)
110            metrics['classification_report'] = class_report
111        return metrics
112
113    def save_model(self, filepath: str):
114        if not self.trained:
115            raise ValueError("Model must be trained before saving")
116        torch.save({
117            'model_state_dict': self.model.state_dict(),
118            'config': self.config,
119            'feature_names': self.feature_names,
120            'class_names': self.class_names,
121            'trained': self.trained
122        }, filepath)
123        print(f"GNN model saved to {filepath}")
124
125    def load_model(self, filepath: str):
126        checkpoint = torch.load(filepath, map_location=self.device)
127        self.model = SimpleGCN(
128            self.config['input_dim'],
129            self.config['hidden_dim'],
130            self.config['output_dim']
131        ).to(self.device)
132        self.model.load_state_dict(checkpoint['model_state_dict'])
133        self.config = checkpoint.get('config', self.config)
134        self.feature_names = checkpoint.get('feature_names', [])
135        self.class_names = checkpoint.get('class_names', [])
136        self.trained = checkpoint.get('trained', True)
137        print(f"GNN model loaded from {filepath}")