Working on test loop

Finished model training and accuracy

Author: rzimmerdev

model.py

@@ -50,7 +50,7 @@ def forward(self, x):
         x = x.permute(1, 0, 2)
         x = x[:, -1, :]
         x = self.classifier(x)
-        return x
+        return nn.functional.softmax(x, dim=1)
 
     def training_step(self, batch, batch_idx):
         x, y = batch
@@ -72,7 +72,9 @@ 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)
+        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)
         return loss

test.py

@@ -1 +1,41 @@
+import lightning
+import torch
+from torch.utils.data import DataLoader
+from model import MultiClassAnomaly
+from dataset import SeriesDataset
 
+
+def test(args):
+    checkpoint = torch.load(args.checkpoint_path)
+    hyperparams = checkpoint['hyperparameters']
+
+    model_args = {
+        'input_size': hyperparams['input_size'],
+        'hidden_size': hyperparams['hidden_size'],
+        'num_heads': hyperparams['num_heads'],
+        'num_layers': hyperparams['num_layers'],
+        'num_classes': hyperparams['num_classes'],
+    }
+
+    model = MultiClassAnomaly(**model_args)
+    model.load_state_dict(checkpoint['state_dict'])
+    model.eval()
+
+    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)
+
+
+if __name__ == '__main__':
+    import argparse
+
+    parser = argparse.ArgumentParser(description='Test trained model on the entire dataset')
+    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')
+
+    args = parser.parse_args()
+    test(args)

train.py

@@ -14,15 +14,22 @@ def train(args):
     dataset = SeriesDataset(args.data_dir)
 
     train_size = int(0.8 * len(dataset))
-
     train_dataset, val_dataset = torch.utils.data.random_split(dataset, [train_size, len(dataset) - train_size])
 
     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=100)
+    trainer = lightning.Trainer(default_root_dir=args.checkpoint_dir, max_epochs=args.epochs)
     trainer.fit(model=model, train_dataloaders=train_dataloader, val_dataloaders=val_dataloader)
 
+    checkpoint = {
+        'hyperparameters': vars(args),
+        'state_dict': model.state_dict(),
+    }
+    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')
@@ -32,12 +39,13 @@ 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('--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')
 
-    train(parser.parse_args())
+    args = parser.parse_args()
+    train(args)
 
 
 if __name__ == '__main__':