[ADD] calc acl and mcl

1. consistency_measure
2. metric
3. new test func
4. fixed val mmnist set
5. logging (should have no prob)

Author: sentient-codebot

.vscode/launch.json

@@ -71,7 +71,11 @@
                 "--decode_hidden",
                 "false",
                 "--num_slots",
-                "3"
+                "3",
+                "--task",
+                "MMNIST",
+                "--use_val_set",
+                "false",
                 // "--batch_size",
                 // "8"
             ],

argument_parser.py

@@ -95,7 +95,6 @@ def _to_int(foo):
         else: #['1', '2', '3']
             return [int(foo) for foo in foo]
     return _to_int(list__)
-print(mmnist_num_obj('1,2,3,4;1;0'))
 
 def argument_parser():
     """Function to parse all the arguments"""
@@ -118,7 +117,7 @@ def argument_parser():
     parser.add_argument('--ball_trainset', type=str2ballset, default=None, help='train set for ball task')
     parser.add_argument('--ball_testset', type=str2ballset, default=None, help='test set for ball task')
 
-    parser.add_argument('--mmnist_num_objects', '--num_objects', '--num_obj', type=mmnist_num_obj, default=[[2],[2],[1,2,3]], 
+    parser.add_argument('--mmnist_num_objects', '--num_objects', '--num_obj', type=mmnist_num_obj, default=[[2],[2],[2]], 
                         help='number of objects in the MMNIST task (train/test/val). default: 2;2;1,2,3')
 
     # Training Settings
@@ -138,6 +137,7 @@ def argument_parser():
     parser.add_argument('--test_frequency', type=int, default=10,
                         metavar="Frequency at which we log the intermediate variables of the model",
                         help='Just type in a positive integer')
+    parser.add_argument('--use_val_set', type=str2bool, default=False)
     parser.add_argument('--path_to_load_model', type=str, default="",
                         metavar='Relative Path to load the model',
                         help='Relative Path to load the model. If this is empty, no model'
@@ -297,7 +297,8 @@ def argument_parser():
     elif args.task == 'VOR':
         args.mot_gt_file = os.path.join(args.dataset_dir, 'gt_jsons', 'vor_test.json')
 
-
+    if args.use_val_set == True and args.task != 'MMNIST':
+        raise NotImplementedError
 
     return args
 

datasets/MovingMNIST.py

@@ -45,11 +45,13 @@ class MovingMNIST(data.Dataset):
         ("train-labels-idx1-ubyte.gz", "d53e105ee54ea40749a09fcbcd1e9432"),
         ("mnist_test_seq.npy", "be083ec986bfe91a449d63653c411eb2"),
     ]
+    val_dataset = 'mmnist_val.pt'
     def __init__(self, root, train=True, n_frames_input=10, n_frames_output=10, num_objects=[2],
                 static_prob=-1,
                 download=False,
                 transform=None,
-                length=int(1e4),):
+                length=int(1e4),
+                val=False):
         '''
         Args:
             `root`: Root directory of the dataset (mnist dataset and moving mnist test set)
@@ -69,6 +71,10 @@ def __init__(self, root, train=True, n_frames_input=10, n_frames_output=10, num_
         super(MovingMNIST, self).__init__()
         self.root = root
         self.is_train = train
+        if not self.is_train:
+            self.is_val = val
+        else:
+            self.is_val = False
 
         if download:
             self.download()
@@ -79,12 +85,22 @@ def __init__(self, root, train=True, n_frames_input=10, n_frames_output=10, num_
         self.dataset = None
         if train:
             self.mnist, self.mnist_label = load_mnist(root)
+        elif self.is_val:
+            if num_objects[0] != 2:
+                self.mnist, self.mnist_label = load_mnist(root)
+            else:
+                self.dataset = torch.load(os.path.join(root, self.val_dataset))
         else:
             if num_objects[0] != 2:
                 self.mnist, self.mnist_label = load_mnist(root)
             else:
                 self.dataset = load_fixed_set(root, False)
-        self.length = length if self.dataset is None else self.dataset.shape[1]
+        if self.dataset is None:
+            self.length = length
+        elif self.is_val:
+            self.length = len(self.dataset)
+        else:
+            self.length = self.dataset.shape[1]
 
         self.num_objects = num_objects
         self.n_frames_input = n_frames_input
@@ -180,6 +196,9 @@ def __getitem__(self, idx):
             num_digits = random.choice(self.num_objects)
             # Generate data on the fly
             images, ind_images, labels = self.generate_moving_mnist(num_digits)
+        elif self.is_val:
+            labels, input, output, ind_images = *self.dataset[idx],
+            return labels, input, output, ind_images
         else:
             images = self.dataset[:, idx, ...]
 

datasets/__init__.py

@@ -74,11 +74,12 @@ def setup_dataloader(args):
         )
         val_set = MovingMNIST(
             root=args.dataset_dir,
-            train=True,
+            train=False,
             n_frames_input=10,
             n_frames_output=10,
             num_objects=args.mmnist_num_objects[2],# 1 2 3
-            download=True
+            download=True,
+            val=True,
         )
     elif args.task == 'BBALL':
         train_set = BouncingBall(root=args.dataset_dir, train=True, length=20, filename=args.ball_trainset)
@@ -158,7 +159,7 @@ def setup_dataloader(args):
             dataset=val_set,
             batch_size=args.batch_size,
             shuffle=True,
-            num_workers=4 if not DEBUG else 0,
+            num_workers=0 if not DEBUG else 0,
             worker_init_fn=seed_worker,
             generator=g,
         )

gen_mmnist.py

@@ -0,0 +1,37 @@
+import torch
+import numpy as np
+from datasets import MovingMNIST
+import argparse
+from tqdm import tqdm
+
+# ind_image = np.load('ind_images.npy')
+
+val_set = MovingMNIST(
+            root='data',
+            train=True,
+            n_frames_input=10,
+            n_frames_output=10,
+            num_objects=[2],# 1 2 3
+            download=False,
+            length=2000
+)
+
+# list_labels, list_input, list_output, list_ind_images = [], [], [], []
+# for idx in tqdm(range(len(val_set))):
+#     labels,input,output, ind_images = val_set[idx]
+#     list_labels.append(labels)
+#     list_input.append(input)
+#     list_output.append(output)
+#     list_ind_images.append(ind_images)
+
+# tensors = torch.stack(list_labels), torch.stack(list_input), torch.stack(list_output), torch.stack(list_ind_images)
+# names = ['labels.npy', 'input.npy', 'output.npy', 'ind_images.npy']
+# for name, tensor in zip(names, tensors):
+#     np.save(name, tensor.numpy())
+
+dataset = []
+for idx in tqdm(range(len(val_set))):
+    dataset.append(val_set[idx])
+    
+torch.save(dataset, 'mmnist_val.pt')
+...

networks.py

@@ -742,7 +742,7 @@ def forward(self, x, h_prev, M_prev=None):
             if "SEP" in self.decoder_type:
                 curr_dec_out_, curr_channels, curr_alpha_mask = self.decoder(encoded_input)
                 next_dec_out_, next_channels, next_alpha_mask = self.decoder(pred_latent)
-                if self.do_logging:
+                if self.do_logging or True: # always log ind_output
                     blocked_out_ = next_channels*next_alpha_mask
                     self.hidden_features['individual_output'] = blocked_out_.detach()
                     self.hidden_features['individual_recons'] = (curr_channels*curr_alpha_mask).detach()
@@ -752,7 +752,7 @@ def forward(self, x, h_prev, M_prev=None):
         else:
             if "SEP" in self.decoder_type:
                 next_dec_out_, next_channels, next_alpha_mask = self.decoder(h_new)
-                if self.do_logging:
+                if self.do_logging or True: # always log ind_output
                     blocked_out_ = next_channels*next_alpha_mask
                     self.hidden_features['individual_output'] = blocked_out_.detach()
             else:

test_mmnist.py

@@ -16,6 +16,7 @@
 from utils.visualize import VecStack, make_grid_video, plot_heatmap, mplfig_to_video
 from utils.logging import log_stats, enable_logging, setup_wandb_columns
 from utils.metric import f1_score, gen_masks, get_mot_metrics
+from utils.metric import consistency_measure
 from tqdm import tqdm
 import wandb
 from utils import util
@@ -46,7 +47,7 @@ def get_grad_norm(model):
     return total_norm
 
 # @torch.no_grad()
-def test(model, test_loader, args, loss_fn, writer, rollout=True, epoch=0, log_columns=None):
+def test(model, test_loader, args, loss_fn, writer, rollout=True, epoch=0, log_columns=None, calc_csty=False):
     '''test(model, test_loader, args, loss_fn, writer, rollout)'''
     start_time = time()
     # wandb table
@@ -90,6 +91,8 @@ def test(model, test_loader, args, loss_fn, writer, rollout=True, epoch=0, log_c
     ssim = 0.
     most_used_units = []
     pred_list = []
+    epoch_avr_len = 0.
+    epoch_max_len = 0.
     id_counter = 0
     for batch_idx, data in enumerate(tqdm(test_loader) if __name__ == "__main__" else test_loader): # tqdm doesn't work here?
         if args.task == 'MMNIST':
@@ -131,6 +134,7 @@ def test(model, test_loader, args, loss_fn, writer, rollout=True, epoch=0, log_c
             data.shape[3],
             data.shape[4])
         ) # (BS, num_blocks, T, C, H, W)
+        ind_pred = torch.empty((data.shape[0], args.num_hidden, data.shape[1]-rollout_start, data.shape[2], data.shape[3], data.shape[4]))
         reconstruction = []
         individual_recons = []
         soft_masks = [] # list of batches of masks
@@ -171,6 +175,8 @@ def test(model, test_loader, args, loss_fn, writer, rollout=True, epoch=0, log_c
             f1 += f1_frame
 
             prediction[:, frame+1, :, :, :] = preds
+            if frame >= rollout_start:
+                ind_pred[:, :, frame-rollout_start, :, :, :] = model.hidden_features['individual_output']
             if do_logging:
                 blocked_prediction[:, 0, frame+1, :, :, :] = preds # dim == 6
                 blocked_prediction[:, 1:, frame+1, :, :, :] = model.hidden_features['individual_output']
@@ -255,6 +261,14 @@ def test(model, test_loader, args, loss_fn, writer, rollout=True, epoch=0, log_c
                 soft_masks=torch.stack(soft_masks, dim=1).cpu(), # [BS, T, K, H, W]
             )
 
+        # calculate consistency
+        avr_len, max_len = None, None
+        if 'SEP' in args.decoder_type and calc_csty:
+            avr_len, max_len = consistency_measure(ind_pred, ind_digits[:, :, rollout_start:, ...], 
+                                                   corr_padding=(1,1), output_ids=False, reduction='mean', exclude_background=True)
+            epoch_avr_len += avr_len
+            epoch_max_len += max_len    
+        
         if not rollout:
             ssim += pt_ssim.ssim(data[:,1:,:,:,:].reshape((-1,1,data.shape[3],data.shape[4])), # data.shape = (batch, frame, 1, height, width)
                         prediction[:,1:,:,:,:].reshape((-1,1,data.shape[3],data.shape[4])))
@@ -284,6 +298,8 @@ def test(model, test_loader, args, loss_fn, writer, rollout=True, epoch=0, log_c
     epoch_recon_loss /= len(test_loader)
     epoch_pred_loss /= len(test_loader)
     epoch_mseloss = epoch_mseloss / (batch_idx+1)
+    epoch_avr_len /= len(test_loader)
+    epoch_max_len /= len(test_loader)
     ssim = ssim / (batch_idx+1)
     f1_avg = f1 / (batch_idx+1) / (data.shape[1]-1)
 
@@ -336,6 +352,9 @@ def test(model, test_loader, args, loss_fn, writer, rollout=True, epoch=0, log_c
         metrics['rule_attn_probs_sm'] = torch.stack(rule_attn_probs_sm, dim=1) # Shape: [N, T, num_hidden, num_rules]
     if len(rule_attn_probs_gsm) > 0:
         metrics['rule_attn_probs_gsm'] = torch.stack(rule_attn_probs_gsm, dim=1) # Shape: [N, T, num_hidden, num_rules]
+    if 'SEP' in args.decoder_type and calc_csty:
+        metrics['avr_len'] = epoch_avr_len
+        metrics['max_len'] = epoch_max_len    
 
     # slot attention
     if args.use_slot_attention:
@@ -346,6 +365,7 @@ def test(model, test_loader, args, loss_fn, writer, rollout=True, epoch=0, log_c
     print('test runtime:', time() - start_time)
     return epoch_loss, epoch_recon_loss, epoch_pred_loss, prediction, data, metrics, test_table
 
+
 @torch.no_grad()
 def dec_rim_util(model, h):
     """check the contribution of the (num_module)-th RIM 
@@ -425,13 +445,14 @@ def main():
     # call test function
     test_loss, recon_loss, pred_loss, prediction, data, metrics, test_table = test(
         model = model,
-        test_loader = test_loader,
+        test_loader = val_loader if args.use_val_set else test_loader,
         args = args,
         loss_fn = loss_fn,
         writer = writer,
         rollout = True,
         epoch = epoch,
         log_columns = columns,
+        calc_csty = True if args.use_val_set else False,
     )
     log_stats(
         args=args,

test_mmnist_val.py

@@ -0,0 +1,12 @@
+from argument_parser import argument_parser
+from datasets import setup_dataloader
+
+
+args = argument_parser()
+
+foo, val_loader, bar = setup_dataloader(args)
+
+print(len(val_loader.dataset))
+t = next(iter(val_loader))
+print(t[3].shape)
+...

train_mmnist.py

@@ -18,7 +18,7 @@
 from utils.logging import log_stats, setup_wandb_columns
 from datasets import setup_dataloader
 from tqdm import tqdm
-from test_mmnist import dec_rim_util, test
+from test_mmnist import test
 
 import os 
 from os import listdir
@@ -123,7 +123,7 @@ def main():
     columns = setup_wandb_columns(args) # artifact columns
 
     # data setup
-    train_loader, _, test_loader = setup_dataloader(args=args)
+    train_loader, val_loader, test_loader = setup_dataloader(args=args)
 
     # model setup
     model, optimizer, scheduler, loss_fn, start_epoch, train_batch_idx, best_mse = setup_model(args=args)
@@ -160,13 +160,14 @@ def main():
             """test model accuracy and log intermediate variables here"""
             test_loss, test_recon_loss, test_pred_loss, prediction, data, metrics, test_table = test(
                 model = model, 
-                test_loader = test_loader, 
+                test_loader = val_loader if args.use_val_set else test_loader, 
                 args = args, 
                 loss_fn = loss_fn, 
                 writer = writer,
                 rollout = True,
                 epoch = epoch,
                 log_columns=columns if epoch%50==0 else None,
+                calc_csty = True if args.use_val_set else False
             )
             log_stats(
                 args=args,

utils/consistensy_measure/__init__.py

@@ -8,24 +8,39 @@ def consistency_measure(
     target_seq: torch.Tensor,
     corr_padding: tuple=(0, 0),
     output_ids: bool=True,
+    reduction: str='none',
+    exclude_background: bool=True,
 ):
     """
     input:
         `input_seq`: [N, K1, T, C, H, W]
         `target_seq`: [N, K2, T, C, H, W]
         `corr_padding`: (h-wise, w-wise) padding for correlation operation
+    return:
+        `avr_len`, `max_len`, (`IDs`)
     """
     input_seq = input_seq.permute(2, 0, 1, 3, 4, 5) # [T, N, K1, C, H, W]
     target_seq = target_seq.permute(2, 0, 1, 3, 4, 5) # [T, N, K2, C, H, W]
     IDs = []
     for t in range(input_seq.shape[0]):
         corr_coef = normalized_corr(input_seq[t], target_seq[t], padding=corr_padding) # [N, K1, K2]
         _, indices = torch.max(corr_coef, dim=-1) # indices, shape [N, K1,]. 
+        bg_flag = input_seq[t].sum(dim=(-1, -2, -3)) < 0.1 * target_seq[t].sum(dim=(-1,-2,-3)).mean(-1, keepdim=True) # [N, K1]
+        indices[bg_flag] = target_seq.shape[1]+1 # extra ID for background
         IDs.append(indices)
     IDs = torch.stack(IDs, dim=-1) # shape [N, 3, T]
     avr_len = average_consistent_length(IDs)
     max_len = maximum_consistent_length(IDs)
 
+    if reduction == 'mean':
+        avr_len = avr_len.mean()
+        max_len = max_len.mean()
+    elif reduction == 'sum':
+        avr_len = avr_len.sum()
+        max_len = max_len.sum()
+    else:
+        pass
+    
     if output_ids:
         return avr_len, max_len, IDs
     else: