Finished test metrics

Author: rzimmerdev

requirements.txt

@@ -78,3 +78,8 @@ urllib3=2.1.0=pypi_0
 wheel=0.41.3=pyhd8ed1ab_0
 xz=5.2.6=h166bdaf_0
 yarl=1.9.2=pypi_0
+lightning~=2.1.1
+torch~=2.1.0
+plotly~=5.18.0
+torchmetrics~=1.2.0
+pandas~=1.5.3

src/model.py

@@ -9,6 +9,8 @@ def __init__(self, input_size, hidden_size, num_heads, num_layers, num_classes,
         super().__init__()
         self.accuracy = torchmetrics.Accuracy(task='multiclass', num_classes=num_classes)
 
+        self.num_classes = num_classes
+
         self.cnn = nn.Sequential(
             nn.Sequential(
                 nn.Conv1d(input_size, hidden_size, kernel_size=3, padding=1),
@@ -75,8 +77,26 @@ def test_step(self, batch, batch_idx):
         pred = torch.argmax(logits, dim=1)
         y = torch.argmax(y, dim=1)
         acc = self.accuracy(pred, y)
-        self.log('test_loss', loss, on_epoch=True, on_step=False)
-        self.log('test_accuracy', acc, on_epoch=True, on_step=False)
+
+        cm = torchmetrics.functional.confusion_matrix(pred, y, task='multiclass', num_classes=self.num_classes)
+        # Can't log tensors, and cm is multiclass, so have to log each class separately
+        self.log('loss', loss, on_epoch=True, on_step=False)
+        self.log('accuracy', acc, on_epoch=True, on_step=False)
+
+        for i in range(self.num_classes):
+            false_positives = torch.sum(cm[:, i]) - cm[i, i]
+            false_negatives = torch.sum(cm[i, :]) - cm[i, i]
+            true_positives = cm[i, i]
+            true_negatives = torch.sum(cm) - (false_positives + false_negatives + true_positives)
+
+            precision = true_positives / (true_positives + false_positives + 1e-8)
+            recall = true_positives / (true_positives + false_negatives + 1e-8)
+            f1 = 2 * (precision * recall) / (precision + recall + 1e-8)
+
+            self.log(f'precision_{i}', precision, on_epoch=True, on_step=False)
+            self.log(f'recall_{i}', recall, on_epoch=True, on_step=False)
+            self.log(f'f1_{i}', f1, on_epoch=True, on_step=False)
+
         return loss
 
     def configure_optimizers(self):

src/split.py

@@ -0,0 +1,30 @@
+import os
+import shutil
+import random
+
+split_ratio = 0.2
+test_dir = "data/test"
+train_dir = "data/train"
+data_dir = "data/Sara_dataset"
+
+os.makedirs(train_dir, exist_ok=True)
+os.makedirs(test_dir, exist_ok=True)
+
+file_list = os.listdir(data_dir)
+
+test_size = int(split_ratio * len(file_list))
+
+random.shuffle(file_list)
+
+for file_name in file_list[:test_size]:
+    source_path = os.path.join(data_dir, file_name)
+    target_path = os.path.join(test_dir, file_name)
+    shutil.move(source_path, target_path)
+
+for file_name in file_list[test_size:]:
+    source_path = os.path.join(data_dir, file_name)
+    target_path = os.path.join(train_dir, file_name)
+    shutil.move(source_path, target_path)
+
+# Remove data dir
+shutil.rmtree(data_dir)

src/test.py

@@ -5,6 +5,9 @@
 from src.dataset import SeriesDataset
 
 
+from plotly import express as px
+
+
 def test(args):
     checkpoint = torch.load(args.checkpoint_path)
     hyperparams = checkpoint['hyperparameters']
@@ -24,8 +27,9 @@ def test(args):
     dataset = SeriesDataset(args.data_dir)
     dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False)
 
-    trainer = lightning.Trainer(default_root_dir=args['checkpoint_dir'], max_epochs=args['epochs'])
-    trainer.test(model=model, dataloaders=dataloader)
+    trainer = lightning.Trainer(default_root_dir=hyperparams['checkpoint_dir'], max_epochs=hyperparams['max_epochs'])
+
+    metrics = trainer.test(model=model, dataloaders=dataloader)
 
 
 if __name__ == '__main__':
@@ -35,7 +39,7 @@ def test(args):
     parser.add_argument('--checkpoint_path', type=str, default="checkpoints/model_checkpoint.pth",
                         help='Path to the model checkpoint file')
     parser.add_argument('--batch_size', type=int, default=8, help='Batch size for testing')
-    parser.add_argument('--data_dir', type=str, default="Sara_dataset/test", help='Dataset directory for loading series')
+    parser.add_argument('--data_dir', type=str, default="data/test", help='Dataset directory for loading series')
 
     args = parser.parse_args()
     test(args)

src/train.py

@@ -19,7 +19,7 @@ def train(args):
     train_dataloader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)
     val_dataloader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False)
 
-    trainer = lightning.Trainer(default_root_dir=args.checkpoint_dir, max_epochs=args.epochs)
+    trainer = lightning.Trainer(default_root_dir=args.checkpoint_dir, max_epochs=args.max_epochs)
     trainer.fit(model=model, train_dataloaders=train_dataloader, val_dataloaders=val_dataloader)
 
     checkpoint = {
@@ -28,8 +28,6 @@ def train(args):
     }
     torch.save(checkpoint, f'{args.checkpoint_dir}/model_checkpoint.pth')
 
-    trainer.test(model=model, dataloaders=val_dataloader)
-
 
 def main():
     parser = argparse.ArgumentParser(description='CNN-Transformer Time Series Classification')
@@ -39,10 +37,10 @@ def main():
     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="Sara_dataset/", help='Dataset directory for loading series')
+    parser.add_argument('--data_dir', type=str, default="data/train", 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')
-    parser.add_argument('--epochs', type=int, default=100, help='Number of epochs to train for')
+    parser.add_argument('--max_epochs', type=int, default=100, help='Number of epochs to train for')
 
     args = parser.parse_args()
     train(args)