├── README.md ├── LICENSE ├── evaluate.py ├── train.py ├── gnn.py └── gsl.py /README.md: -------------------------------------------------------------------------------- 1 | # Graph Neural Networks with Trainable Adjacency Matrices for Fault Diagnosis on Multivariate Sensor Data 2 | 3 | This repository is the official implementation of model architectures from the paper [Graph Neural Networks with Trainable Adjacency Matrices for Fault Diagnosis on Multivariate Sensor Data](https://doi.org/10.1109/ACCESS.2024.3481331). 4 | 5 | ## Training and evaluation examples 6 | 7 | [FDDBenchmark](https://github.com/AIRI-Institute/fddbenchmark) was used in our experiments. 8 | 9 | Training step: 10 | 11 | ``` 12 | python train.py 13 | ``` 14 | 15 | Evaluation step: 16 | 17 | ``` 18 | python evaluate.py 19 | ``` 20 | 21 | ## Citation 22 | 23 | Please cite our paper as follows: 24 | 25 | ``` 26 | @article{kovalenko2024graph, 27 | title={Graph neural networks with trainable adjacency matrices for fault diagnosis on multivariate sensor data}, 28 | author={Kovalenko, Aleksandr and Pozdnyakov, Vitaliy and Makarov, Ilya}, 29 | journal={IEEE Access}, 30 | year={2024}, 31 | volume={12}, 32 | pages={152860-152872}, 33 | publisher={IEEE} 34 | } 35 | ``` 36 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2024 AI Research 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /evaluate.py: -------------------------------------------------------------------------------- 1 | import pandas as pd 2 | from sklearn.preprocessing import StandardScaler 3 | import torch 4 | import argparse 5 | 6 | from fddbenchmark import FDDDataset, FDDDataloader, FDDEvaluator 7 | from gnn import GNN_TAM 8 | 9 | 10 | def parse_args(): 11 | parser = argparse.ArgumentParser(description='model_inference') 12 | parser.add_argument('--dataset', type=str, default='reinartz_tep') 13 | parser.add_argument('--window_size', type=int, default=100) 14 | parser.add_argument('--step_size', type=int, default=1) 15 | parser.add_argument('--batch_size', type=int, default=512) 16 | parser.add_argument('--name', type=str, default='gnn1') 17 | return parser.parse_args() 18 | 19 | 20 | def inference(): 21 | args = parse_args() 22 | device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 23 | print('Using device:', device) 24 | # Data preparation: 25 | dataset = FDDDataset(name=args.dataset) 26 | scaler = StandardScaler() 27 | scaler.fit(dataset.df[dataset.train_mask]) 28 | dataset.df[:] = scaler.transform(dataset.df) 29 | test_dl = FDDDataloader( 30 | dataframe=dataset.df, 31 | label=dataset.label, 32 | mask=dataset.test_mask, 33 | window_size=args.window_size, 34 | step_size=args.step_size, 35 | use_minibatches=True, 36 | batch_size=args.batch_size, 37 | shuffle=True 38 | ) 39 | # Load saved model: 40 | model = torch.load('saved_models/' + args.name + '.pt', 41 | map_location=device) 42 | # Inference: 43 | model.eval() 44 | preds = [] 45 | test_labels = [] 46 | for test_ts, test_index, test_label in test_dl: 47 | ts = torch.FloatTensor(test_ts).to(device) 48 | ts = torch.transpose(ts, 1, 2) 49 | with torch.no_grad(): 50 | logits = model(ts) 51 | pred = logits.argmax(axis=1).cpu().numpy() 52 | preds.append(pd.Series(pred, index=test_index)) 53 | test_labels.append(pd.Series(test_label, index=test_index)) 54 | pred = pd.concat(preds) 55 | test_label = pd.concat(test_labels) 56 | 57 | evaluator = FDDEvaluator(step_size=1) 58 | evaluator.print_metrics(test_label, pred) 59 | 60 | 61 | if __name__ == '__main__': 62 | inference() 63 | -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | from sklearn.preprocessing import StandardScaler 2 | import torch 3 | import torch.nn.functional as F 4 | from torch.optim import Adam 5 | from tqdm.auto import tqdm, trange 6 | import argparse 7 | 8 | from fddbenchmark import FDDDataset, FDDDataloader 9 | from gnn import GNN_TAM 10 | 11 | 12 | def parse_args(): 13 | parser = argparse.ArgumentParser(description='train model') 14 | parser.add_argument('--dataset', type=str, default='reinartz_tep') 15 | parser.add_argument('--n_epochs', type=int, default=40) 16 | parser.add_argument('--window_size', type=int, default=100) 17 | parser.add_argument('--step_size', type=int, default=1) 18 | parser.add_argument('--batch_size', type=int, default=512) 19 | parser.add_argument('--n_gnn', type=int, default=1) 20 | parser.add_argument('--gsl_type', type=str, default='tanh') 21 | parser.add_argument('--n_hidden', type=int, default=1024) 22 | parser.add_argument('--alpha', type=float, default=0.1) 23 | parser.add_argument('--k', type=int, default=None) 24 | parser.add_argument('--name', type=str, default='gnn') 25 | return parser.parse_args() 26 | 27 | 28 | def train(): 29 | args = parse_args() 30 | device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 31 | print('Using device:', device) 32 | # Data preparation: 33 | dataset = FDDDataset(name=args.dataset) 34 | scaler = StandardScaler() 35 | scaler.fit(dataset.df[dataset.train_mask]) 36 | dataset.df[:] = scaler.transform(dataset.df) 37 | n_nodes = dataset.df.shape[1] 38 | n_classes = len(set(dataset.label)) 39 | train_dl = FDDDataloader( 40 | dataframe=dataset.df, 41 | label=dataset.label, 42 | mask=dataset.train_mask, 43 | window_size=args.window_size, 44 | step_size=args.step_size, 45 | use_minibatches=True, 46 | batch_size=args.batch_size, 47 | shuffle=True 48 | ) 49 | # Model creation: 50 | model = GNN_TAM(n_nodes=n_nodes, 51 | window_size=args.window_size, 52 | n_classes=n_classes, 53 | n_gnn=args.n_gnn, 54 | gsl_type=args.gsl_type, 55 | n_hidden=args.n_hidden, 56 | alpha=args.alpha, 57 | k=args.k, 58 | device=device) 59 | model.to(device) 60 | # Training: 61 | model.train() 62 | optimizer = Adam(model.parameters(), lr=0.001) 63 | weight = torch.ones(n_classes) * 0.5 64 | weight[1:] /= (n_classes - 1) 65 | for e in trange(args.n_epochs, desc="Epochs ..."): 66 | av_loss = [] 67 | for train_ts, train_index, train_label in tqdm(train_dl): 68 | ts = torch.FloatTensor(train_ts).to(device) 69 | ts = torch.transpose(ts, 1, 2) 70 | train_label = torch.LongTensor(train_label).to(device) 71 | logits = model(ts) 72 | loss = F.cross_entropy(logits, train_label, weight=weight.to(device)) 73 | optimizer.zero_grad() 74 | loss.backward() 75 | optimizer.step() 76 | av_loss.append(loss.item()) 77 | print(f'Epoch: {e+1:2d}/{args.n_epochs}, average CE loss: {sum(av_loss)/len(av_loss):.4f}') 78 | 79 | torch.save(model, 'saved_models/' + args.name + str(args.n_gnn) + '.pt') 80 | 81 | 82 | if __name__ == '__main__': 83 | train() 84 | -------------------------------------------------------------------------------- /gnn.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | from gsl import GSL 5 | 6 | 7 | class GCLayer(nn.Module): 8 | """ 9 | Graph convolution layer. 10 | """ 11 | def __init__(self, in_dim, out_dim): 12 | super().__init__() 13 | self.dense = nn.Linear(in_dim, out_dim) 14 | 15 | def forward(self, adj, X): 16 | adj = adj + torch.eye(adj.size(0)).to(adj.device) 17 | h = self.dense(X) 18 | norm = adj.sum(1)**(-1/2) 19 | h = norm[None, :] * adj * norm[:, None] @ h 20 | return h 21 | 22 | 23 | class GNN_TAM(nn.Module): 24 | """ 25 | Model architecture from the paper "Graph Neural Networks with Trainable 26 | Adjacency Matrices for Fault Diagnosis on Multivariate Sensor Data". 27 | https://doi.org/10.1109/ACCESS.2024.3481331 28 | """ 29 | def __init__( 30 | self, 31 | n_nodes: int, 32 | window_size: int, 33 | n_classes: int, 34 | n_gnn: int = 1, 35 | gsl_type: str = 'relu', 36 | n_hidden: int = 1024, 37 | alpha: float = 0.1, 38 | k: int = None, 39 | device: str = 'cpu' 40 | ): 41 | """ 42 | Args: 43 | n_nodes (int): The number of nodes/sensors. 44 | window_size (int): The number of timestamps in one sample. 45 | n_classes (int): The number of classes. 46 | n_gnn (int): The number of GNN modules. 47 | gsl_type (str): The type of GSL block. 48 | n_hidden (int): The number of hidden parameters in GCN layers. 49 | alpha (float): Saturation rate for GSL block. 50 | k (int): The maximum number of edges from one node. 51 | device (str): The name of a device to train the model. `cpu` and 52 | `cuda` are possible. 53 | """ 54 | super(GNN_TAM, self).__init__() 55 | self.window_size = window_size 56 | self.nhidden = n_hidden 57 | self.device = device 58 | self.idx = torch.arange(n_nodes).to(device) 59 | self.adj = [0 for i in range(n_gnn)] 60 | self.h = [0 for i in range(n_gnn)] 61 | self.skip = [0 for i in range(n_gnn)] 62 | self.z = (torch.ones(n_nodes, n_nodes) - torch.eye(n_nodes)).to(device) 63 | self.n_gnn = n_gnn 64 | 65 | self.gsl = nn.ModuleList() 66 | self.conv1 = nn.ModuleList() 67 | self.bnorm1 = nn.ModuleList() 68 | self.conv2 = nn.ModuleList() 69 | self.bnorm2 = nn.ModuleList() 70 | 71 | for i in range(self.n_gnn): 72 | self.gsl.append(GSL(gsl_type, n_nodes, 73 | window_size, alpha, k, device)) 74 | self.conv1.append(GCLayer(window_size, n_hidden)) 75 | self.bnorm1.append(nn.BatchNorm1d(n_nodes)) 76 | self.conv2.append(GCLayer(n_hidden, n_hidden)) 77 | self.bnorm2.append(nn.BatchNorm1d(n_nodes)) 78 | 79 | self.fc = nn.Linear(n_gnn*n_hidden, n_classes) 80 | 81 | def forward(self, X): 82 | X = X.to(self.device) 83 | for i in range(self.n_gnn): 84 | self.adj[i] = self.gsl[i](self.idx) 85 | self.adj[i] = self.adj[i] * self.z 86 | self.h[i] = self.conv1[i](self.adj[i], X).relu() 87 | self.h[i] = self.bnorm1[i](self.h[i]) 88 | self.skip[i], _ = torch.min(self.h[i], dim=1) 89 | self.h[i] = self.conv2[i](self.adj[i], self.h[i]).relu() 90 | self.h[i] = self.bnorm2[i](self.h[i]) 91 | self.h[i], _ = torch.min(self.h[i], dim=1) 92 | self.h[i] = self.h[i] + self.skip[i] 93 | 94 | h = torch.cat(self.h, 1) 95 | output = self.fc(h) 96 | 97 | return output 98 | 99 | def get_adj(self): 100 | return self.adj 101 | -------------------------------------------------------------------------------- /gsl.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | 6 | # A = ReLu(W) 7 | class Graph_ReLu_W(nn.Module): 8 | def __init__(self, n_nodes, k, device): 9 | super(Graph_ReLu_W, self).__init__() 10 | self.num_nodes = n_nodes 11 | self.k = k 12 | self.A = nn.Parameter(torch.randn(n_nodes, n_nodes).to(device), 13 | requires_grad=True).to(device) 14 | 15 | def forward(self, idx): 16 | adj = F.relu(self.A) 17 | if self.k: 18 | mask = torch.zeros(idx.size(0), idx.size(0)).to(self.device) 19 | mask.fill_(float('0')) 20 | v, id = (adj + torch.rand_like(adj)*0.01).topk(self.k, 1) 21 | mask.scatter_(1, id, v.fill_(1)) 22 | adj = adj*mask 23 | return adj 24 | 25 | 26 | # A for Directed graphs: 27 | class Graph_Directed_A(nn.Module): 28 | def __init__(self, n_nodes, window_size, alpha, k, device): 29 | super(Graph_Directed_A, self).__init__() 30 | self.alpha = alpha 31 | self.k = k 32 | self.device = device 33 | self.e1 = nn.Embedding(n_nodes, window_size) 34 | self.e2 = nn.Embedding(n_nodes, window_size) 35 | self.l1 = nn.Linear(window_size, window_size) 36 | self.l2 = nn.Linear(window_size, window_size) 37 | 38 | def forward(self, idx): 39 | m1 = torch.tanh(self.alpha*self.l1(self.e1(idx))) 40 | m2 = torch.tanh(self.alpha*self.l2(self.e2(idx))) 41 | adj = F.relu(torch.tanh(self.alpha*torch.mm(m1, m2.transpose(1, 0)))) 42 | if self.k: 43 | mask = torch.zeros(idx.size(0), idx.size(0)).to(self.device) 44 | mask.fill_(float('0')) 45 | v, id = (adj + torch.rand_like(adj)*0.01).topk(self.k, 1) 46 | mask.scatter_(1, id, v.fill_(1)) 47 | adj = adj*mask 48 | return adj 49 | 50 | 51 | # A for Uni-directed graphs: 52 | class Graph_Uni_Directed_A(nn.Module): 53 | def __init__(self, n_nodes, window_size, alpha, k, device): 54 | super(Graph_Directed_A, self).__init__() 55 | self.alpha = alpha 56 | self.k = k 57 | self.device = device 58 | self.e1 = nn.Embedding(n_nodes, window_size) 59 | self.e2 = nn.Embedding(n_nodes, window_size) 60 | self.l1 = nn.Linear(window_size, window_size) 61 | self.l2 = nn.Linear(window_size, window_size) 62 | 63 | def forward(self, idx): 64 | m1 = torch.tanh(self.alpha*self.l1(self.e1(idx))) 65 | m2 = torch.tanh(self.alpha*self.l2(self.e2(idx))) 66 | adj = F.relu(torch.tanh(self.alpha*(torch.mm(m1, m2.transpose(1, 0)) 67 | - torch.mm(m2, m1.transpose(1, 0))))) 68 | if self.k: 69 | mask = torch.zeros(idx.size(0), idx.size(0)).to(self.device) 70 | mask.fill_(float('0')) 71 | v, id = (adj + torch.rand_like(adj)*0.01).topk(self.k, 1) 72 | mask.scatter_(1, id, v.fill_(1)) 73 | adj = adj*mask 74 | return adj 75 | 76 | 77 | # A for Undirected graphs: 78 | class Graph_Undirected_A(nn.Module): 79 | def __init__(self, n_nodes, window_size, alpha, k, device): 80 | super(Graph_Directed_A, self).__init__() 81 | self.alpha = alpha 82 | self.k = k 83 | self.device = device 84 | self.e1 = nn.Embedding(n_nodes, window_size) 85 | self.l1 = nn.Linear(window_size, window_size) 86 | 87 | def forward(self, idx): 88 | m1 = torch.tanh(self.alpha*self.l1(self.e1(idx))) 89 | m2 = torch.tanh(self.alpha*self.l1(self.e1(idx))) 90 | adj = F.relu(torch.tanh(self.alpha*torch.mm(m1, m2.transpose(1, 0)))) 91 | if self.k: 92 | mask = torch.zeros(idx.size(0), idx.size(0)).to(self.device) 93 | mask.fill_(float('0')) 94 | v, id = (adj + torch.rand_like(adj)*0.01).topk(self.k, 1) 95 | mask.scatter_(1, id, v.fill_(1)) 96 | adj = adj*mask 97 | return adj 98 | 99 | 100 | class GSL(nn.Module): 101 | """ 102 | Graph structure learning block. 103 | """ 104 | def __init__( 105 | self, 106 | gsl_type, 107 | n_nodes, 108 | window_size, 109 | alpha, 110 | k, 111 | device): 112 | super(GSL, self).__init__() 113 | self.gsl_layer = None 114 | if gsl_type == 'relu': 115 | self.gsl_layer = Graph_ReLu_W(n_nodes, k, device) 116 | elif gsl_type == 'directed': 117 | self.self.gsl_layer = Graph_Directed_A(n_nodes, window_size, 118 | alpha, k, device) 119 | elif gsl_type == 'unidirected': 120 | self.self.gsl_layer = Graph_Uni_Directed_A(n_nodes, window_size, 121 | alpha, k, device) 122 | elif gsl_type == 'undirected': 123 | self.self.gsl_layer = Graph_Undirected_A(n_nodes, window_size, 124 | alpha, k, device) 125 | else: 126 | print('Wrong name of graph structure learning layer!') 127 | 128 | def forward(self, idx): 129 | return self.gsl_layer(idx) 130 | --------------------------------------------------------------------------------