Initial commit

Author: rzimmerdev

.gitignore

@@ -0,0 +1,3 @@
+checkpoints/
+lightning_logs/
+Sara_dataset/

README.md

@@ -1,2 +1,8 @@
 # ml-anomaly
 Time Series Multi-class Anomaly detection using a Hybrid Transformer/CNN for physical sensor data readings. Provides an improvement on the existing dataset analysis paper Evaluating Conveyor Belt Health With Signal Processing Applied to Inertial Sensing
+
+## Installation
+```bash
+pip install -r requirements.txt
+```
+

dataset.py

@@ -1,9 +1,28 @@
 import os
+from enum import Enum
+
 import pandas as pd
+import torch
 from torch.utils.data import Dataset
 
 
+class Labels(Enum):
+    NORMAL = 0
+    ANOMALY_1 = 1
+    ANOMALY_2 = 2
+    ANOMALY_3 = 3
+    ANOMALY_4 = 4
+
+
 class SeriesDataset(Dataset):
+    label_map = {
+        "Situação referencia index": 0,
+        "Situação rolo 3 da esquerda levantado index": 1,
+        "Situação rolo 3 da esquerda removido index": 2,
+        "Situação rolo 3 da direita levantado index": 3,
+        "Situação rolo 3 da direita removido index": 4,
+    }
+
     def __init__(self, data_dir):
         self.data_dir = data_dir
         self.file_list = [f for f in os.listdir(data_dir) if f.endswith('.csv')]
@@ -12,9 +31,13 @@ def __len__(self):
         return len(self.file_list)
 
     def __getitem__(self, idx):
+        # filename = <LABEL> <ID>.csv
         file_path = os.path.join(self.data_dir, self.file_list[idx])
-        data = pd.read_csv(file_path).values  # Assuming CSV files have numeric values
-        x = torch.FloatTensor(data[:, :-1])  # Input features
-        y = torch.LongTensor(data[:, -1])  # Labels
-        return x, y
 
+        label = " ".join(self.file_list[idx].split('.')[0].split(' ')[:-1])
+        y = torch.zeros(5).scatter_(0, torch.LongTensor([self.label_map[label]]), 1)
+
+        data = pd.read_csv(file_path).values
+        x = torch.FloatTensor(data)
+
+        return x, y

model.py

@@ -1,36 +1,81 @@
 import torch
-import torch.nn as nn
+import torchmetrics
+from torch import nn
+import lightning
 
-class MultiClassAnomaly(nn.Module):
-    def __init__(self, input_size, hidden_size, num_heads, num_layers, num_classes, dropout_rate=0.1):
-        super(MultiClassAnomaly, self).__init__()
+
+class MultiClassAnomaly(lightning.LightningModule):
+    def __init__(self, input_size, hidden_size, num_heads, num_layers, num_classes, dropout_rate=0.5):
+        super().__init__()
+        self.accuracy = torchmetrics.Accuracy(task='multiclass', num_classes=num_classes)
 
         self.cnn = nn.Sequential(
-            nn.Conv1d(in_channels=input_size, out_channels=hidden_size, kernel_size=3, padding=1),
-            nn.ReLU(),
-            nn.MaxPool1d(kernel_size=2)
+            nn.Sequential(
+                nn.Conv1d(input_size, hidden_size, kernel_size=3, padding=1),
+                nn.BatchNorm1d(hidden_size),
+                nn.ReLU(),
+                nn.MaxPool1d(2),
+            ),
+            nn.Sequential(
+                nn.Conv1d(hidden_size, hidden_size, kernel_size=3, padding=1),
+                nn.BatchNorm1d(hidden_size),
+                nn.ReLU(),
+                nn.MaxPool1d(2),
+            )
         )
 
-        self.embedding = nn.Linear(hidden_size, hidden_size)
-        self.transformer_encoder = nn.TransformerEncoder(
-            nn.TransformerEncoderLayer(d_model=hidden_size, nhead=num_heads),
+        self.transformer = nn.TransformerEncoder(
+            nn.TransformerEncoderLayer(
+                d_model=hidden_size,
+                nhead=num_heads,
+                dim_feedforward=hidden_size,
+                dropout=dropout_rate,
+                activation='relu'
+            ),
             num_layers=num_layers
         )
 
-        self.fc = nn.Linear(hidden_size, num_classes)
-
-        self.dropout = nn.Dropout(p=dropout_rate)
+        self.classifier = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size),
+            nn.ReLU(),
+            nn.Dropout(dropout_rate),
+            nn.Linear(hidden_size, num_classes)
+        )
 
     def forward(self, x):
-        x = self.cnn(x)
         x = x.permute(0, 2, 1)
-        x = self.embedding(x)
+        x = self.cnn(x)
+        x = x.permute(2, 0, 1)
+        x = self.transformer(x)
+        x = x.permute(1, 0, 2)
+        x = x[:, -1, :]
+        x = self.classifier(x)
+        return x
 
-        x = self.transformer_encoder(x)
-        x = x.mean(dim=1)
+    def training_step(self, batch, batch_idx):
+        x, y = batch
+        logits = self.forward(x)
+        loss = nn.functional.cross_entropy(logits, y)
+        self.log('train_loss', loss, on_epoch=True, on_step=True)
+        return loss
 
-        x = self.dropout(x)
-        output = self.fc(x)
+    def validation_step(self, batch, batch_idx):
+        x, y = batch
+        logits = self.forward(x)
+        loss = nn.functional.cross_entropy(logits, y)
+        acc = self.accuracy(logits, y)
+        self.log('val_loss', loss, on_epoch=True, on_step=False)
+        self.log('val_accuracy', acc, on_epoch=True, on_step=False)
+        return loss
 
-        return output
+    def test_step(self, batch, batch_idx):
+        x, y = batch
+        logits = self.forward(x)
+        loss = nn.functional.cross_entropy(logits, y)
+        acc = self.accuracy(logits, y)
+        self.log('test_loss', loss, on_epoch=True, on_step=False)
+        self.log('test_accuracy', acc, on_epoch=True, on_step=False)
+        return loss
 
+    def configure_optimizers(self):
+        return torch.optim.Adam(self.parameters(), lr=0.001)

requirements.txt

@@ -0,0 +1,80 @@
+# This file may be used to create an environment using:
+# $ conda create --name <env> --file <this file>
+# platform: linux-64
+_libgcc_mutex=0.1=conda_forge
+_openmp_mutex=4.5=2_gnu
+aiohttp=3.8.6=pypi_0
+aiosignal=1.3.1=pypi_0
+async-timeout=4.0.3=pypi_0
+attrs=23.1.0=pypi_0
+bzip2=1.0.8=hd590300_5
+ca-certificates=2023.7.22=hbcca054_0
+certifi=2023.7.22=pypi_0
+charset-normalizer=3.3.2=pypi_0
+filelock=3.13.1=pypi_0
+frozenlist=1.4.0=pypi_0
+fsspec=2023.10.0=pypi_0
+idna=3.4=pypi_0
+jinja2=3.1.2=pypi_0
+ld_impl_linux-64=2.40=h41732ed_0
+libblas=3.9.0=19_linux64_openblas
+libcblas=3.9.0=19_linux64_openblas
+libffi=3.4.2=h7f98852_5
+libgcc-ng=13.2.0=h807b86a_3
+libgfortran-ng=13.2.0=h69a702a_3
+libgfortran5=13.2.0=ha4646dd_3
+libgomp=13.2.0=h807b86a_3
+liblapack=3.9.0=19_linux64_openblas
+libnsl=2.0.1=hd590300_0
+libopenblas=0.3.24=pthreads_h413a1c8_0
+libsqlite=3.44.0=h2797004_0
+libstdcxx-ng=13.2.0=h7e041cc_3
+libuuid=2.38.1=h0b41bf4_0
+libzlib=1.2.13=hd590300_5
+lightning=2.1.1=pypi_0
+lightning-utilities=0.9.0=pypi_0
+markupsafe=2.1.3=pypi_0
+mpmath=1.3.0=pypi_0
+multidict=6.0.4=pypi_0
+ncurses=6.4=h59595ed_2
+networkx=3.2.1=pypi_0
+numpy=1.26.2=pypi_0
+nvidia-cublas-cu12=12.1.3.1=pypi_0
+nvidia-cuda-cupti-cu12=12.1.105=pypi_0
+nvidia-cuda-nvrtc-cu12=12.1.105=pypi_0
+nvidia-cuda-runtime-cu12=12.1.105=pypi_0
+nvidia-cudnn-cu12=8.9.2.26=pypi_0
+nvidia-cufft-cu12=11.0.2.54=pypi_0
+nvidia-curand-cu12=10.3.2.106=pypi_0
+nvidia-cusolver-cu12=11.4.5.107=pypi_0
+nvidia-cusparse-cu12=12.1.0.106=pypi_0
+nvidia-nccl-cu12=2.18.1=pypi_0
+nvidia-nvjitlink-cu12=12.3.52=pypi_0
+nvidia-nvtx-cu12=12.1.105=pypi_0
+openssl=3.1.4=hd590300_0
+packaging=23.2=pypi_0
+pandas=2.1.3=py310hcc13569_0
+pip=23.3.1=pyhd8ed1ab_0
+python=3.10.13=hd12c33a_0_cpython
+python-dateutil=2.8.2=pyhd8ed1ab_0
+python-tzdata=2023.3=pyhd8ed1ab_0
+python_abi=3.10=4_cp310
+pytorch-lightning=2.1.1=pypi_0
+pytz=2023.3.post1=pyhd8ed1ab_0
+pyyaml=6.0.1=pypi_0
+readline=8.2=h8228510_1
+requests=2.31.0=pypi_0
+setuptools=68.2.2=pyhd8ed1ab_0
+six=1.16.0=pyh6c4a22f_0
+sympy=1.12=pypi_0
+tk=8.6.13=noxft_h4845f30_101
+torch=2.1.0=pypi_0
+torchmetrics=1.2.0=pypi_0
+tqdm=4.66.1=pypi_0
+triton=2.1.0=pypi_0
+typing-extensions=4.8.0=pypi_0
+tzdata=2023c=h71feb2d_0
+urllib3=2.1.0=pypi_0
+wheel=0.41.3=pyhd8ed1ab_0
+xz=5.2.6=h166bdaf_0
+yarl=1.9.2=pypi_0

train.py

@@ -1,48 +1,44 @@
+import argparse
 
+import torch
+from torch.utils.data import DataLoader
+import lightning
 
-def train_model(args):
-    model = CNNTransformerModel(args.input_size, args.hidden_size, args.num_heads, args.num_layers, args.num_classes)
+from model import MultiClassAnomaly
+from dataset import SeriesDataset
 
-    model_checkpoint = pl.callbacks.ModelCheckpoint(
-        dirpath=args.checkpoint_dir,
-        filename='best_model',
-        monitor='val_loss',
-        mode='min',
-        save_top_k=1
-    )
 
-    train_dataset = TimeSeriesDataset(args.data_dir)
-    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)
+def train(args):
+    model = MultiClassAnomaly(args.input_size, args.hidden_size, args.num_heads, args.num_layers, args.num_classes)
 
-    val_dataset = TimeSeriesDataset(args.val_data_dir)
-    val_loader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False)
+    dataset = SeriesDataset(args.data_dir)
 
-    trainer = CNNTransformerTrainer(
-        model,
-        train_loader=train_loader,
-        val_loader=val_loader,
-        test_loader=test_loader,
-        model_checkpoint=model_checkpoint
-    )
+    train_size = int(0.8 * len(dataset))
 
-    if args.resume_training:
-        trainer.load_checkpoint(args.resume_checkpoint)
+    train_dataset, val_dataset = torch.utils.data.random_split(dataset, [train_size, len(dataset) - train_size])
 
-    trainer.fit()
+    train_dataloader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)
+    val_dataloader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False)
 
-    path = os.path.join(args.checkpoint_dir, 'final_model.pth')
-    trainer.save_checkpoint(path)
+    trainer = lightning.Trainer(default_root_dir=args.checkpoint_dir, max_epochs=100)
+    trainer.fit(model=model, train_dataloaders=train_dataloader, val_dataloaders=val_dataloader)
 
 
 def main():
     parser = argparse.ArgumentParser(description='CNN-Transformer Time Series Classification')
-    parser.add_argument('--input_size', type=int, default=6, help='Number of input features')
+    parser.add_argument('--input_size', type=int, default=7, help='Number of input features')
     parser.add_argument('--hidden_size', type=int, default=64, help='Hidden size for the model')
     parser.add_argument('--num_heads', type=int, default=4, help='Number of attention heads')
     parser.add_argument('--num_layers', type=int, default=2, help='Number of transformer layers')
     parser.add_argument('--num_classes', type=int, default=5, help='Number of classes for classification')
     parser.add_argument('--batch_size', type=int, default=8, help='Batch size for training')
-    parser.add_argument('--data_dir', type=str, default, help='Dataset directory from which to load series')
 
-    train(parser.args)
+    parser.add_argument('--data_dir', type=str, default="Sara_dataset/", help='Dataset directory for loading series')
+    parser.add_argument('--checkpoint_dir', type=str, default="checkpoints/", help='Directory to save checkpoints')
+    parser.add_argument('--resume_training', action='store_true', help='Resume training from checkpoint')
 
+    train(parser.parse_args())
+
+
+if __name__ == '__main__':
+    main()