feat: added tensorboard to classifier and multitask classifier

.gitignore

@@ -1,3 +1,4 @@
 .idea
 __pycache__
-sst-classifier.pt
+sst-classifier.pt
+logdir

classifier.py

@@ -1,4 +1,5 @@
 import time, random, numpy as np, argparse, sys, re, os
+from datetime import datetime
 from types import SimpleNamespace
 import csv
 
@@ -7,6 +8,7 @@
 from torch import nn
 from torch.utils.data import Dataset, DataLoader
 from sklearn.metrics import classification_report, f1_score, recall_score, accuracy_score
+from torch.utils.tensorboard import SummaryWriter
 
 # change it with respect to the original model
 from tokenizer import BertTokenizer
@@ -232,6 +234,10 @@ def save_model(model, optimizer, args, config, filepath):
 
 
 def train(args):
+    name = datetime.now().strftime("%Y%m%d-%H%M%S")
+    writer = SummaryWriter(log_dir=args.logdir + "/classifier/" + name)
+    loss_idx_value = 0
+
     device = torch.device('cuda') if args.use_gpu else torch.device('cpu')
     # Load data
     # Create the data and its corresponding datasets and dataloader
@@ -283,19 +289,28 @@ def train(args):
             optimizer.step()
 
             train_loss += loss.item()
+            writer.add_scalar("Loss/Minibatches", loss.item(), loss_idx_value)
+            loss_idx_value += 1
             num_batches += 1
 
         train_loss = train_loss / (num_batches)
+        writer.add_scalar("Loss/Epochs", train_loss, epoch)
 
         train_acc, train_f1, *_ = model_eval(train_dataloader, model, device)
+        writer.add_scalar("Accuracy/train/Epochs", train_acc, epoch)
+        writer.add_scalar("F1_score/train/Epochs", train_f1, epoch)
+
         dev_acc, dev_f1, *_ = model_eval(dev_dataloader, model, device)
+        writer.add_scalar("Accuracy/dev/Epochs", dev_acc, epoch)
+        writer.add_scalar("F1_score/dev/Epochs", dev_f1, epoch)
 
         if dev_acc > best_dev_acc:
             best_dev_acc = dev_acc
             save_model(model, optimizer, args, config, args.filepath)
 
         print(
             f"Epoch {epoch}: train loss :: {train_loss :.3f}, train acc :: {train_acc :.3f}, dev acc :: {dev_acc :.3f}")
+    writer.close()
 
 
 def test(args):
@@ -345,6 +360,8 @@ def get_args():
     parser.add_argument("--dev_out", type=str, default="cfimdb-dev-output.txt")
     parser.add_argument("--test_out", type=str, default="cfimdb-test-output.txt")
 
+    parser.add_argument("--logdir", type=str, default="logdir")
+
     parser.add_argument("--batch_size", help='sst: 64, cfimdb: 8 can fit a 12GB GPU', type=int, default=8)
     parser.add_argument("--hidden_dropout_prob", type=float, default=0.3)
     parser.add_argument("--lr", type=float, help="learning rate, default lr for 'pretrain': 1e-3, 'finetune': 1e-5",
@@ -372,7 +389,8 @@ def get_args():
         test='data/ids-sst-test-student.csv',
         option=args.option,
         dev_out='predictions/' + args.option + '-sst-dev-out.csv',
-        test_out='predictions/' + args.option + '-sst-test-out.csv'
+        test_out='predictions/' + args.option + '-sst-test-out.csv',
+        logdir=args.logdir
     )
 
     train(config)

multitask_classifier.py

@@ -1,10 +1,12 @@
 import time, random, numpy as np, argparse, sys, re, os
+from datetime import datetime
 from types import SimpleNamespace
 
 import torch
 from torch import nn
 import torch.nn.functional as F
 from torch.utils.data import DataLoader
+from torch.utils.tensorboard import SummaryWriter
 
 from bert import BertModel
 from optimizer import AdamW
@@ -109,6 +111,10 @@ def save_model(model, optimizer, args, config, filepath):
 
 ## Currently only trains on sst dataset
 def train_multitask(args):
+    name = datetime.now().strftime("%Y%m%d-%H%M%S")
+    writer = SummaryWriter(log_dir=args.logdir + "/multitask_classifier/" + name)
+    loss_idx_value = 0
+
     device = torch.device('cuda') if args.use_gpu else torch.device('cpu')
     # Load data
     # Create the data and its corresponding datasets and dataloader
@@ -185,6 +191,7 @@ def train_multitask(args):
             sts_loss = F.nll_loss(logits, b_labels.view(-1))
 
             train_loss += sts_loss.item()
+            writer.add_scalar("Loss/STS/Minibatches", sts_loss.item(), loss_idx_value)
             num_batches += 1
 
             # Train on PARAPHRASE dataset
@@ -206,6 +213,7 @@ def train_multitask(args):
             para_loss = F.nll_loss(logits, b_labels.view(-1))
 
             train_loss += para_loss.item()
+            writer.add_scalar("Loss/PARA/Minibatches", para_loss.item(), loss_idx_value)
             num_batches += 1
 
             # Train on SST dataset
@@ -222,6 +230,8 @@ def train_multitask(args):
             sst_loss = F.cross_entropy(logits, b_labels.view(-1))
 
             train_loss += sst_loss.item()
+            writer.add_scalar("Loss/SST/Minibatches", sst_loss.item(), loss_idx_value)
+            loss_idx_value += 1
             num_batches += 1
 
             # Calculate gradient and update weights
@@ -231,15 +241,23 @@ def train_multitask(args):
             optimizer.step()
 
         train_loss = train_loss / num_batches
+        writer.add_scalar("Loss/Epochs", train_loss, epoch)
 
         para_train_acc, _, _, sst_train_acc, _, _, sts_train_acc, _, _ = model_eval_multitask(sst_train_dataloader,
                                                                                               para_train_dataloader,
                                                                                               sts_train_dataloader,
                                                                                               model, device)
+        writer.add_scalar("para_acc/train/Epochs", para_train_acc, epoch)
+        writer.add_scalar("sst_acc/train/Epochs", sst_train_acc, epoch)
+        writer.add_scalar("sts_acc/train/Epochs", sts_train_acc, epoch)
+
         para_dev_acc, _, _, sst_dev_acc, _, _, sts_dev_acc, _, _ = model_eval_multitask(sst_dev_dataloader,
                                                                                         para_dev_dataloader,
                                                                                         sts_dev_dataloader, model,
                                                                                         device)
+        writer.add_scalar("para_acc/dev/Epochs", para_dev_acc, epoch)
+        writer.add_scalar("sst_acc/dev/Epochs", sst_dev_acc, epoch)
+        writer.add_scalar("sts_acc/dev/Epochs", sts_dev_acc, epoch)
 
         if para_dev_acc > best_dev_acc_para and sst_dev_acc > best_dev_acc_sst and sts_dev_acc > best_dev_acc_sts:
             best_dev_acc_para = para_dev_acc
@@ -248,6 +266,9 @@ def train_multitask(args):
             save_model(model, optimizer, args, config, args.filepath)
         train_acc = sst_train_acc + para_train_acc + sts_train_acc
         dev_acc = sst_dev_acc + para_dev_acc + sts_dev_acc
+
+        writer.add_scalar("acc/train/Epochs", train_acc, epoch)
+        writer.add_scalar("acc/dev/Epochs", dev_acc, epoch)
         print(
             f"Epoch {epoch}: train loss :: {train_loss :.3f}, combined train acc :: {train_acc :.3f}, combined dev acc :: {dev_acc :.3f}")
 
@@ -296,6 +317,8 @@ def get_args():
     parser.add_argument("--sts_dev_out", type=str, default="predictions/sts-dev-output.csv")
     parser.add_argument("--sts_test_out", type=str, default="predictions/sts-test-output.csv")
 
+    parser.add_argument("--logdir", type=str, default="logdir")
+
     # hyper parameters
     parser.add_argument("--batch_size", help='sst: 64, cfimdb: 8 can fit a 12GB GPU', type=int, default=8)
     parser.add_argument("--hidden_dropout_prob", type=float, default=0.3)

setup.sh

@@ -11,3 +11,4 @@ pip install filelock==3.0.12
 pip install sklearn==0.0
 pip install tokenizers==0.10.1
 pip install explainaboard_client==0.0.7
+pip install tensorboard