add evaluation code during training

.gitignore

@@ -0,0 +1,7 @@
+.vscode
+*.h5
+*.pth
+log_dir
+save_dir
+score_dir
+*.pyc

config.py

@@ -1,5 +1,3 @@
-# -*- coding: utf-8 -*-
-
 import pprint
 
 
@@ -8,7 +6,7 @@ class Config():
     def __init__(self, **kwargs):
 
         # Path
-        self.data_path = 'fcsn_TVSum.h5'
+        self.data_path = '../data/fcsn_tvsum.h5'
         self.save_dir = 'save_dir'
         self.score_dir = 'score_dir'
         self.log_dir = 'log_dir'
@@ -19,6 +17,7 @@ def __init__(self, **kwargs):
         self.n_epochs = 50
         self.n_class = 2
         self.lr = 1e-3
+        self.momentum = 0.9
         self.batch_size = 5
 
         for k, v in kwargs.items():

data_loader.py

@@ -1,5 +1,3 @@
-# -*- coding: utf-8 -*-
-
 import torch
 from torch.utils.data import DataLoader
 
@@ -24,12 +22,10 @@ def __getitem__(self, index):
 
 def get_loader(path, batch_size=5):
     dataset = VideoData(path)
-    train_dataset, test_dataset = torch.utils.data.random_split(dataset, [9*len(dataset)//10, len(dataset)//10])
+    train_dataset, test_dataset = torch.utils.data.random_split(dataset, [len(dataset) - len(dataset)//10, len(dataset)//10])
     train_loader = DataLoader(train_dataset, batch_size=batch_size)
     return train_loader, test_dataset
 
 
 if __name__ == '__main__':
     loader = get_loader('fcsn_dataset.h5')
-    import ipdb
-    ipdb.set_trace()

fcsn.py

@@ -1,5 +1,3 @@
-# -*- coding: utf-8 -*-
-
 import torch.nn as nn
 from collections import OrderedDict 
 

knapsack.py

@@ -0,0 +1,66 @@
+# http://www.geeksforgeeks.org/knapsack-problem/
+
+import numpy as np
+
+def knapsack(v, w, max_weight):
+    rows = len(v) + 1
+    cols = max_weight + 1
+
+    # adding dummy values as later on we consider these values as indexed from 1 for convinence
+
+    v = np.r_[[0], v]
+    w = np.r_[[0], w]
+
+    # row : values , #col : weights
+    dp_array = [[0 for i in range(cols)] for j in range(rows)]
+
+    # 0th row and 0th column have value 0
+
+    # values
+    for i in range(1, rows):
+        # weights
+        for j in range(1, cols):
+            # if this weight exceeds max_weight at that point
+            if j - w[i] < 0:
+                dp_array[i][j] = dp_array[i - 1][j]
+
+            # max of -> last ele taken | this ele taken + max of previous values possible
+            else:
+                dp_array[i][j] = max(dp_array[i - 1][j], v[i] + dp_array[i - 1][j - w[i]])
+
+    # return dp_array[rows][cols]  : will have the max value possible for given wieghts
+
+    chosen = []
+    i = rows - 1
+    j = cols - 1
+
+    # Get the items to be picked
+    while i > 0 and j > 0:
+
+        # ith element is added
+        if dp_array[i][j] != dp_array[i - 1][j]:
+            # add the value
+            chosen.append(i-1)
+            # decrease the weight possible (j)
+            j = j - w[i]
+            # go to previous row
+            i = i - 1
+
+        else:
+            i = i - 1
+
+    return dp_array[rows - 1][cols - 1], chosen
+
+# main
+if __name__ == "__main__":
+    values = list(map(int, input().split()))
+    weights = list(map(int, input().split()))
+    max_weight = int(input())
+
+    max_value, chosen = knapsack(values, weights, max_weight)
+
+    print("The max value possible is")
+    print(max_value)
+
+    print("The index chosen for these are")
+    print(' '.join(str(x) for x in chosen))

tools.py

@@ -0,0 +1,59 @@
+import torch
+import numpy as np
+
+from knapsack import knapsack
+
+
+def eval_metrics(y_pred: torch.Tensor, y_true: torch.Tensor):
+    overlap = (y_pred * y_true).sum().item()
+    precision = overlap / (y_pred.sum().item() + 1e-8)
+    recall = overlap / (y_true.sum().item() + 1e-8)
+    if precision == 0 and recall == 0:
+        fscore = 0
+    else:
+        fscore = 2 * precision * recall / (precision + recall)
+    return [precision, recall, fscore]
+
+
+def select_keyshots(num_frames, cps, weight, value):
+    _, selected = knapsack(value, weight, int(0.15 * num_frames))
+    selected = selected[::-1]
+    key_labels = np.zeros(shape=(num_frames,))
+    for i in selected:
+        key_labels[cps[i][0]:cps[i][1]] = 1
+    return selected, key_labels
+
+
+def upsample(down_arr, N):
+    up_arr = np.zeros(N)
+    ratio = N // 320
+    l = (N - ratio * 320) // 2
+    i = 0
+    while i < 320:
+        up_arr[l:l+ratio] = np.ones(ratio, dtype=int) * down_arr[i]
+        l += ratio
+        i += 1
+    return up_arr
+
+
+def eval_single(video_info, pred_score):
+    """
+    Evaluate F-score of given video and pred_score.
+
+    Args:
+        video_info: hdf5 dataset instance, containing necessary infomation for evaluation.
+        pred_score: output of FCSN model.
+    Returns:
+        evaluation result (precision, recall, f-score).
+    """
+    N = video_info['length'][()]
+    cps = video_info['change_points'][()]
+    weight = video_info['n_frame_per_seg'][()]
+    true_summary_arr = video_info['user_summary'][()]
+    pred_score = np.array(pred_score.cpu().data)
+    pred_score = upsample(pred_score, N)
+    pred_value = np.array([pred_score[cp[0]:cp[1]].mean() for cp in cps])
+    pred_selected, pred_summary = select_keyshots(N, cps, weight, pred_value)
+    eval_arr = [eval_metrics(pred_summary, true_summary) for true_summary in true_summary_arr]
+    eval_res = np.mean(eval_arr, axis=0)
+    return eval_res.tolist()

train.py

@@ -1,39 +1,45 @@
-# -*- coding: utf-8 -*-
-
 import torch
 import torch.optim as optim
+
 from tensorboardX import SummaryWriter
 import numpy as np
 import json
 import os
 from tqdm import tqdm, trange
+import h5py
+from prettytable import PrettyTable
 
 from fcsn import FCSN
-
+import tools
 
 class Solver(object):
     """Class that Builds, Trains FCSN model"""
 
-    def __init__(self, config=None, train_loader=None, test_loader=None):
+    def __init__(self, config=None, train_loader=None, test_dataset=None):
         self.config = config
         self.train_loader = train_loader
-        self.test_loader = test_loader
+        self.test_dataset = test_dataset
 
         # model
         self.model = FCSN(self.config.n_class)
 
         # optimizer
         if self.config.mode == 'train':
             self.optimizer = optim.Adam(self.model.parameters())
-            # self.optimizer = optim.SGD(self.model.parameters(), lr=config.lr, momentum=0.9)
+            # self.optimizer = optim.SGD(self.model.parameters(), lr=config.lr, momentum=self.config.momentum)
             self.model.train()
 
-        # weight
-        self.tvsum_weight = torch.tensor([0.55989996, 4.67362574])
-
         if self.config.gpu:
             self.model = self.model.cuda()
-            self.tvsum_weight = self.tvsum_weight.cuda()
+
+        if not os.path.exists(self.config.score_dir):
+            os.mkdir(self.config.score_dir)
+
+        if not os.path.exists(self.config.save_dir):
+            os.mkdir(self.config.save_dir)
+
+        if not os.path.exists(self.config.log_dir):
+            os.mkdir(self.config.log_dir)
 
     @staticmethod
     def sum_loss(pred_score, gt_labels, weight=None):
@@ -45,11 +51,12 @@ def sum_loss(pred_score, gt_labels, weight=None):
         return loss
 
     def train(self):
-        writer = SummaryWriter()
-        for epoch_i in trange(self.config.n_epochs, desc='Epoch', ncols=80):
+        writer = SummaryWriter(log_dir=self.config.log_dir)
+        t = trange(self.config.n_epochs, desc='Epoch', ncols=80)
+        for epoch_i in t:
             sum_loss_history = []
 
-            for batch_i, (feature, label, _) in enumerate(tqdm(self.train_loader, desc='Batch', ncols=80, leave=False)):
+            for batch_i, (feature, label,  _) in enumerate(tqdm(self.train_loader, desc='Batch', ncols=80, leave=False)):
 
                 # [batch_size, 1024, seq_len]
                 feature.requires_grad_()
@@ -61,25 +68,22 @@ def train(self):
                 # ---- Train ---- #
                 pred_score = self.model(feature)
 
-                # pred_label = torch.argmax(pred_score, dim=1)
-                # loss = torch.nn.MSELoss(pred_label, label)
-                loss = self.sum_loss(pred_score, label, self.tvsum_weight)
+                label_1 = label.sum() / label.shape[0]
+                label_0 = label.shape[1] - label_1
+                weight = torch.tensor([label_1, label_0], dtype=torch.float)
+
+                loss = self.sum_loss(pred_score, label, weight)
                 loss.backward()
 
                 self.optimizer.step()
                 self.optimizer.zero_grad()
                 sum_loss_history.append(loss)
 
             mean_loss = torch.stack(sum_loss_history).mean().item()
-            tqdm.write('\nEpoch {}'.format(epoch_i))
-            tqdm.write('sum loss: {:.3f}'.format(mean_loss))
+            t.set_postfix(loss=mean_loss)
             writer.add_scalar('Loss', mean_loss, epoch_i)
 
-            if (epoch_i+1) % 10 == 0:
-
-                if not os.path.exists(self.config.save_dir):
-                    os.mkdir(self.config.save_dir)
-
+            if (epoch_i+1) % 5 == 0:
                 ckpt_path = self.config.save_dir + '/epoch-{}.pkl'.format(epoch_i)
                 tqdm.write('Save parameters at {}'.format(ckpt_path))
                 torch.save(self.model.state_dict(), ckpt_path)
@@ -89,31 +93,41 @@ def train(self):
     def evaluate(self, epoch_i):
         self.model.eval()
         out_dict = {}
+        eval_arr = []
+        table = PrettyTable()
+        table.title = 'F-score of epoch {}'.format(epoch_i)
+        table.field_names = ['ID', 'Precision', 'Recall', 'F-score']
+        table.float_format = '1.3'
+
+        with h5py.File(self.config.data_path) as data_file:
+            for feature, label, idx in tqdm(self.test_dataset, desc='Evaluate', ncols=80, leave=False):
+                if self.config.gpu:
+                    feature = feature.cuda()
+                pred_score = self.model(feature.unsqueeze(0)).squeeze(0)
+                pred_score = torch.softmax(pred_score, dim=0)[1]
+                video_info = data_file['video_'+str(idx)]
+                eval_res = tools.eval_single(video_info, pred_score)
 
-        for feature, _, idx in tqdm(self.test_loader, desc='Evaluate', ncols=80, leave=False):
-
-            if self.config.gpu:
-                feature = feature.cuda()
-            pred_score = self.model(feature.unsqueeze(0))
-            pred_label = torch.argmax(pred_score, dim=1).squeeze(0).type(dtype=torch.int)
-            pred_label = np.array(pred_label.cpu().data).tolist()
-
-            out_dict[idx] = pred_label
-
-        if not os.path.exists(self.config.score_dir):
-            os.mkdir(self.config.score_dir)
+                eval_arr.append(eval_res)
+                table.add_row([idx] + eval_res)
 
+                pred_score = np.array(pred_score.cpu().data).tolist()
+                out_dict[idx] = pred_score
+
         score_save_path = self.config.score_dir + '/epoch-{}.json'.format(epoch_i)
         with open(score_save_path, 'w') as f:
             tqdm.write('Saving score at {}.'.format(str(score_save_path)))
             json.dump(out_dict, f)
+        eval_mean = np.mean(eval_arr, axis=0).tolist()
+        table.add_row(['mean']+eval_mean)
+        tqdm.write(str(table))
 
 
 if __name__ == '__main__':
     from config import Config
     from data_loader import get_loader
     train_config = Config()
     test_config = Config(mode='test')
-    train_loader, test_loader = get_loader(train_config.data_path, batch_size=train_config.batch_size)
-    solver = Solver(train_config, train_loader, test_loader)
+    train_loader, test_dataset = get_loader(train_config.data_path, batch_size=train_config.batch_size)
+    solver = Solver(train_config, train_loader, test_dataset)
     solver.train()