first commit

watch_and_help/watch_strategy_full/helper.py

@@ -0,0 +1,203 @@
+import os
+import math
+import numpy as np
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.parameter import Parameter
+from torch.nn.modules.rnn import RNNCellBase
+
+def to_cpu(list_of_tensor):
+    if isinstance(list_of_tensor[0], list):
+        list_list_of_tensor = list_of_tensor
+        list_of_tensor = [to_cpu(list_of_tensor)
+                          for list_of_tensor in list_list_of_tensor]
+    else:
+        list_of_tensor = [tensor.cpu() for tensor in list_of_tensor]
+    return list_of_tensor
+
+
+def average_over_list(l):
+    return sum(l) / len(l)
+
+def _LayerNormGRUCell(input, hidden, w_ih, w_hh, ln, b_ih=None, b_hh=None):
+    gi = F.linear(input, w_ih, b_ih)
+    gh = F.linear(hidden, w_hh, b_hh)
+    i_r, i_i, i_n = gi.chunk(3, 1)
+    h_r, h_i, h_n = gh.chunk(3, 1)
+
+    # use layernorm here
+    resetgate = torch.sigmoid(ln['resetgate'](i_r + h_r))
+    inputgate = torch.sigmoid(ln['inputgate'](i_i + h_i))
+    newgate = torch.tanh(ln['newgate'](i_n + resetgate * h_n))
+    hy = newgate + inputgate * (hidden - newgate)
+    return hy
+
+class CombinedEmbedding(nn.Module):
+    def __init__(self, pretrained_embedding, embedding):
+        super(CombinedEmbedding, self).__init__()
+        self.pretrained_embedding = pretrained_embedding
+        self.embedding = embedding
+        self.pivot = pretrained_embedding.num_embeddings
+
+    def forward(self, input):
+        outputs = []
+        mask = input < self.pivot
+        outputs.append(self.pretrained_embedding(torch.clamp(input, 0, self.pivot-1)) * mask.unsqueeze(1).float())
+        mask = input >= self.pivot
+        outputs.append(self.embedding(torch.clamp(input, self.pivot) - self.pivot) * mask.unsqueeze(1).float())
+        return sum(outputs)
+
+
+class writer_helper(object):
+    def __init__(self, writer):
+        self.writer = writer
+        self.all_steps = {}
+
+    def get_step(self, tag):
+        if tag not in self.all_steps.keys():
+            self.all_steps.update({tag: 0})
+
+        step = self.all_steps[tag]
+        self.all_steps[tag] += 1
+        return step
+
+    def scalar_summary(self, tag, value, step=None):
+        if step is None:
+            step = self.get_step(tag)
+        self.writer.add_scalar(tag, value, step)
+
+    def text_summary(self, tag, value, step=None):
+        if step is None:
+            step = self.get_step(tag)
+        self.writer.add_text(tag, value, step)
+
+
+class Constant():
+    def __init__(self, v):
+        self.v = v
+
+    def update(self):
+        pass
+
+
+class LinearStep():
+    def __init__(self, max, min, steps):
+        self.steps = float(steps)
+        self.max = max
+        self.min = min
+        self.cur_step = 0
+        self.v = self.max
+
+    def update(self):
+        v = max(self.max - (self.max - self.min) *
+                self.cur_step / self.steps, self.min)
+        self.cur_step += 1
+        self.v = v
+
+
+class fc_block(nn.Module):
+    def __init__(self, in_channels, out_channels, norm, activation_fn):
+        super(fc_block, self).__init__()
+
+        block = nn.Sequential()
+        block.add_module('linear', nn.Linear(in_channels, out_channels))
+        if norm:
+            block.add_module('batchnorm', nn.BatchNorm1d(out_channels))
+        if activation_fn is not None:
+            block.add_module('activation', activation_fn())
+
+        self.block = block
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class conv_block(nn.Module):
+    def __init__(
+            self,
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride,
+            norm,
+            activation_fn):
+        super(conv_block, self).__init__()
+
+        block = nn.Sequential()
+        block.add_module(
+            'conv',
+            nn.Conv2d(
+                in_channels,
+                out_channels,
+                kernel_size,
+                stride))
+        if norm:
+            block.add_module('batchnorm', nn.BatchNorm2d(out_channels))
+        if activation_fn is not None:
+            block.add_module('activation', activation_fn())
+
+        self.block = block
+
+    def forward(self, x):
+        return self.block(x)
+
+
+def get_conv_output_shape(shape, block):
+    B = 1
+    input = torch.rand(B, *shape)
+    output = block(input)
+    n_size = output.data.view(B, -1).size(1)
+    return n_size
+
+
+class Flatten(nn.Module):
+    def forward(self, input):
+        return input.view(input.size(0), -1)
+
+
+def BHWC_to_BCHW(tensor):
+    tensor = torch.transpose(tensor, 1, 3)  # BCWH
+    tensor = torch.transpose(tensor, 2, 3)  # BCHW
+    return tensor
+
+
+def LCS(X, Y):
+
+    # find the length of the strings
+    m = len(X)
+    n = len(Y)
+
+    # declaring the array for storing the dp values
+    L = [[None] * (n + 1) for i in range(m + 1)]
+    longest_L = [[[]] * (n + 1) for i in range(m + 1)]
+    longest = 0
+    lcs_set = []
+
+    for i in range(m + 1):
+        for j in range(n + 1):
+            if i == 0 or j == 0:
+                L[i][j] = 0
+                longest_L[i][j] = []
+            elif X[i - 1] == Y[j - 1]:
+                L[i][j] = L[i - 1][j - 1] + 1
+                longest_L[i][j] = longest_L[i - 1][j - 1] + [X[i - 1]]
+                if L[i][j] > longest:
+                    lcs_set = []
+                    lcs_set.append(longest_L[i][j])
+                    longest = L[i][j]
+                elif L[i][j] == longest and longest != 0:
+                    lcs_set.append(longest_L[i][j])
+            else:
+                if L[i - 1][j] > L[i][j - 1]:
+                    L[i][j] = L[i - 1][j]
+                    longest_L[i][j] = longest_L[i - 1][j]
+                else:
+                    L[i][j] = L[i][j - 1]
+                    longest_L[i][j] = longest_L[i][j - 1]
+
+    if len(lcs_set) > 0:
+        return lcs_set[0]
+    else:
+        return lcs_set

watch_and_help/watch_strategy_full/network/encoder_decoder.py

@@ -0,0 +1,131 @@
+import copy
+import numpy as np
+from termcolor import colored
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+#import ipdb
+#import pdb
+
+def _align_tensor_index(reference_index, tensor_index):
+
+    where_in_tensor = []
+    for i in reference_index:
+        where = np.where(i == tensor_index)[0][0]
+        where_in_tensor.append(where)
+    return np.array(where_in_tensor)
+
+
+def _sort_by_length(list_of_tensor, batch_length, return_idx=False):    
+    idx = np.argsort(np.array(copy.copy(batch_length)))[::-1]
+    for i, tensor in enumerate(list_of_tensor):
+        if isinstance(tensor, dict):
+            list_of_tensor[i]['class_objects'] = [tensor['class_objects'][j] for j in idx]
+            list_of_tensor[i]['object_coords'] = [tensor['object_coords'][j] for j in idx]
+            list_of_tensor[i]['states_objects'] = [tensor['states_objects'][j] for j in idx]
+            list_of_tensor[i]['mask_object'] = [tensor['mask_object'][j] for j in idx]
+        else:
+            list_of_tensor[i] = [tensor[j] for j in idx]
+    if return_idx:
+        return list_of_tensor, idx
+    else:
+        return list_of_tensor
+
+def _sort_by_index(list_of_tensor, idx):
+    for i, tensor in enumerate(list_of_tensor):
+        list_of_tensor[i] = [tensor[j] for j in idx]
+    return list_of_tensor
+
+class ActionDemo2Predicate(nn.Module):
+    summary_keys = ['loss', 'top1']
+    def __init__(self, args, dset, loss_weight, **kwargs):
+        from network.module_graph import PredicateClassifier
+        super(ActionDemo2Predicate, self).__init__()
+
+        print('------------------------------------------------------------------------------------------')
+        print('ActionDemo2Predicate')
+        print('------------------------------------------------------------------------------------------')
+
+        model_type = kwargs["model_type"]
+        print('model_type', model_type)
+
+        if model_type.lower() == 'max':
+            from network.module_graph import ActionDemoEncoder
+            demo_encoder = ActionDemoEncoder(args, dset,  'max')
+        elif model_type.lower() == 'avg':
+            from network.module_graph import ActionDemoEncoder
+            demo_encoder = ActionDemoEncoder(args, dset, 'avg')
+        elif model_type.lower() == 'lstmavg':
+            from network.module_graph import ActionDemoEncoder
+            demo_encoder = ActionDemoEncoder(args, dset, 'lstmavg')
+        elif model_type.lower() == 'bilstmavg':
+            from network.module_graph import ActionDemoEncoder
+            demo_encoder = ActionDemoEncoder(args, dset, 'bilstmavg')
+        elif model_type.lower() == 'lstmlast':
+            from network.module_graph import ActionDemoEncoder
+            demo_encoder = ActionDemoEncoder(args, dset, 'lstmlast')
+        elif model_type.lower() == 'bilstmlast':
+            from network.module_graph import ActionDemoEncoder
+            demo_encoder = ActionDemoEncoder(args, dset, 'bilstmlast')
+        else:
+            raise ValueError
+        demo_encoder = torch.nn.DataParallel(demo_encoder)
+
+        predicate_decoder = PredicateClassifier(args, dset, loss_weight)
+        
+        # for quick save and load
+        all_modules = nn.Sequential()
+        all_modules.add_module('demo_encoder', demo_encoder)
+        all_modules.add_module('predicate_decoder', predicate_decoder)
+
+        self.demo_encoder = demo_encoder
+        self.predicate_decoder = predicate_decoder
+        self.all_modules = all_modules
+        self.to_cuda_fn = None
+
+    def set_to_cuda_fn(self, to_cuda_fn):
+        self.to_cuda_fn = to_cuda_fn
+
+    def forward(self, data, **kwargs):
+        if self.to_cuda_fn:
+            data = self.to_cuda_fn(data)
+        
+        # demonstration                
+        batch_data = data[0].cuda()
+        batch_gt = data[1].cuda()
+        batch_task_name = data[2]
+        batch_action_id = data[3]
+        
+        # demonstration encoder
+        batch_demo_emb, _ = self.demo_encoder(batch_data, batch_gt, batch_task_name)
+
+        loss, info = self.predicate_decoder(batch_demo_emb, batch_gt, batch_action_id, batch_task_name)
+
+        return loss, info
+                    
+    def write_summary(self, writer, info, postfix):
+
+        model_name = 'Demo2Predicate-{}/'.format(postfix)
+        for k in self.summary_keys:
+            if k in info.keys():
+                writer.scalar_summary(model_name + k, info[k])
+
+    def save(self, path, verbose=False):
+        
+        if verbose:
+            print(colored('[*] Save model at {}'.format(path), 'magenta'))
+        torch.save(self.all_modules.state_dict(), path)
+
+    def load(self, path, verbose=False):
+
+        if verbose:
+            print(colored('[*] Load model at {}'.format(path), 'magenta'))
+        self.all_modules.load_state_dict(
+            torch.load(
+                path,
+                map_location=lambda storage,
+                loc: storage))
+
+
+

watch_and_help/watch_strategy_full/network/module_graph.py

@@ -0,0 +1,249 @@
+import random
+import itertools
+import numpy as np
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.utils.rnn import pad_sequence
+from helper import fc_block, Constant
+
+def _calculate_accuracy_predicate(logits, batch_target, max_possible_count=None, topk=1, multi_classifier=False):
+    batch_size = batch_target.size(0) / max_possible_count
+
+    _, pred = logits.topk(topk, 1, True, True)
+    pred = pred.t()
+    correct = pred.eq(batch_target.view(1, -1).expand_as(pred))
+
+    k = 1    
+    accuray = correct[:k].view(-1).float()
+    accuray = accuray.view(-1, max_possible_count)
+
+    correct_k = (accuray.sum(1)==max_possible_count).sum(0)
+    correct_k = correct_k * (100.0 / batch_size)
+    
+    return correct_k
+
+def _calculate_accuracy(
+        action_correct,
+        object_correct,
+        rel_correct,
+        target_correct,
+        batch_length,
+        info):
+    
+    action_valid_correct = [sum(action_correct[i, :(l - 1)])
+                            for i, l in enumerate(batch_length)]
+    object_valid_correct = [sum(object_correct[i, :(l - 1)])
+                             for i, l in enumerate(batch_length)]
+    rel_valid_correct = [sum(rel_correct[i, :(l - 1)])
+                             for i, l in enumerate(batch_length)]
+    target_valid_correct = [sum(target_correct[i, :(l - 1)])
+                             for i, l in enumerate(batch_length)]
+
+    action_accuracy = sum(action_valid_correct).float() / (sum(batch_length) - 1. * len(batch_length))
+    object_accuracy = sum(object_valid_correct).float() / (sum(batch_length) - 1. * len(batch_length))
+    rel_accuracy = sum(rel_valid_correct).float() / (sum(batch_length) - 1. * len(batch_length))
+    target_accuracy = sum(target_valid_correct).float() / (sum(batch_length) - 1. * len(batch_length))
+
+    info.update({'action_accuracy': action_accuracy.cpu().item()})
+    info.update({'object_accuracy': object_accuracy.cpu().item()})
+    info.update({'rel_accuracy': rel_accuracy.cpu().item()})
+    info.update({'target_accuracy': target_accuracy.cpu().item()})
+
+class PredicateClassifier(nn.Module):
+
+    def __init__(
+            self, 
+            args,
+            dset,
+            loss_weight
+            ):
+
+        super(PredicateClassifier, self).__init__()
+
+        self.num_goal_predicates = dset.num_goal_predicates
+        self.max_possible_count = dset.max_goal_length
+        self.loss = custom_loss(loss_weight)
+
+        hidden_size = args.demo_hidden
+        self.hidden_size = hidden_size
+        self.loss_type = args.loss_type
+        
+        if args.dropout==0:
+            print('dropout', args.dropout)
+            classifier = nn.Sequential()
+            classifier.add_module('fc_block1', fc_block(hidden_size*82, hidden_size, False, nn.Tanh))
+            classifier.add_module('fc_block2', fc_block(hidden_size, 79, False, None)) # 79 is all possible actions
+        else:
+            print('dropout not 0', args.dropout)
+            classifier = nn.Sequential()
+            classifier.add_module('fc_block1', fc_block(hidden_size*82, hidden_size, False, nn.Tanh))
+            classifier.add_module('dropout', nn.Dropout(args.dropout))
+            classifier.add_module('fc_block2', fc_block(hidden_size, 79, False, None)) # 79 is all possible actions
+
+        self.classifier = classifier
+
+    def forward(self, input_emb, batch_target, batch_action_id, batch_task_name, **kwargs):
+        input_emb = input_emb.view(-1, self.hidden_size*82)
+        logits = self.classifier(input_emb)
+
+        prob = F.softmax(logits, 1)
+
+        #cross entropy loss
+        if self.loss_type == 'ce':
+            loss = F.cross_entropy(logits, batch_target)  
+        #custom loss
+        if self.loss_type == 'regu':
+            loss = self.loss(logits, batch_target)  
+
+        argmax_Y = torch.max(logits, 1)[1].view(-1, 1)
+        top1 = (batch_target.float().view(-1, 1) == argmax_Y.float()).sum().item() / len(batch_target.float().view(-1, 1)) * 100
+
+        with torch.no_grad():
+            info = {
+                "prob": prob.cpu().numpy(),
+                "argmax": argmax_Y.cpu().numpy(),
+                "loss": loss.cpu().numpy(),
+                "top1": top1,
+                "target": batch_target.cpu().numpy(),
+                "task_name": batch_task_name,
+                "action_id": batch_action_id.cpu().numpy()
+            }
+
+        return loss, info
+
+class custom_loss(nn.Module):
+	def __init__(self, loss_weight) -> None:
+		super(custom_loss, self).__init__( )
+		self.loss = nn.CrossEntropyLoss().cuda()
+		self.weight_loss = nn.MSELoss().cuda()		
+		self.loss_weight = loss_weight
+		self.counts = torch.FloatTensor(self.loss_weight).cuda()
+
+	def forward(self, pred, target):		
+		# weight loss + cross entropy loss
+		batch_counts = torch.bincount(target)
+		batch_counts = batch_counts/torch.sum(batch_counts)
+		if len(batch_counts) < 79:
+			batch_counts = F.pad(input=batch_counts, pad=(0, 79 - len(batch_counts)%79), mode='constant', value=0)
+		
+		celoss = self.loss(pred, target)
+		customloss = self.weight_loss(batch_counts, self.counts)
+		print('celoss: ', celoss, 'customloss: ', customloss)
+		loss = celoss +  1000*customloss
+		return loss
+
+class PredicateClassifierMultiClassifier(nn.Module):
+
+    def __init__(
+            self, 
+            args,
+            dset):
+
+        super(PredicateClassifierMultiClassifier, self).__init__()
+
+        self.num_goal_predicates = dset.num_goal_predicates
+        self.max_possible_count = dset.max_goal_length
+        self.max_subgoal_length = dset.max_subgoal_length
+
+        hidden_size = args.demo_hidden
+        print('hidden_size', hidden_size)
+        print('PredicateClassifierMultiClassifier')
+
+        if args.dropout==0:
+            print('dropout', args.dropout)
+            classifier = nn.Sequential()
+            classifier.add_module('fc_block1', fc_block(hidden_size, hidden_size, False, nn.Tanh))
+            classifier.add_module('fc_block2', fc_block(hidden_size, self.num_goal_predicates*(self.max_subgoal_length+1), False, None))
+        else:
+            print('dropout not 0', args.dropout)
+            classifier = nn.Sequential()
+            classifier.add_module('fc_block1', fc_block(hidden_size, hidden_size, False, nn.Tanh))
+            classifier.add_module('dropout', nn.Dropout(args.dropout))
+            classifier.add_module('fc_block2', fc_block(hidden_size, self.num_goal_predicates*(self.max_subgoal_length+1), False, None))
+
+        
+        self.classifier = classifier
+
+    def forward(self, bs, input_emb, batch_target, batch_file_name, **kwargs):
+        
+        logits = self.classifier(input_emb)
+        logits = logits.reshape([-1, (self.max_subgoal_length+1)])
+        prob = F.softmax(logits, 1)
+
+        batch_target = torch.cat(batch_target)
+        loss = F.cross_entropy(logits, batch_target)  
+
+        top1 = _calculate_accuracy_predicate(logits, batch_target, self.num_goal_predicates, multi_classifier=True)
+
+        with torch.no_grad():
+            info = {
+                "prob": prob.cpu().numpy(),
+                "loss": loss.cpu().numpy(),
+                "top1": top1.cpu().numpy(),
+                "target": batch_target.cpu().numpy(),
+                "file_name": batch_file_name
+            }
+
+        return loss, info
+
+class ActionDemoEncoder(nn.Module):
+    def __init__(self, args, dset, pooling):
+        super(ActionDemoEncoder, self).__init__()
+        hidden_size = args.demo_hidden
+        self.hidden_size = hidden_size
+
+        len_action_predicates = dset.max_action_len
+        self.action_embed = nn.Embedding(len_action_predicates, hidden_size)
+
+        feat2hidden = nn.Sequential()
+        feat2hidden.add_module(
+            'fc_block1', fc_block(hidden_size, hidden_size, False, nn.ReLU))
+        self.feat2hidden = feat2hidden
+
+        self.pooling = pooling
+
+        if 'lstm' in self.pooling:
+            self.lstm = nn.LSTM(hidden_size, hidden_size)
+
+
+    def forward(self, batch_data, batch_gt, batch_task_name):        
+        batch_data = batch_data.view(-1,1)
+
+        stacked_demo_feat = self.action_embed(batch_data)
+
+        stacked_demo_feat = self.feat2hidden(stacked_demo_feat)
+        batch_demo_feat = []
+        start = 0
+        
+        for length in range(0,batch_data.shape[0]):
+            if length == 0:
+                feat = stacked_demo_feat[0:1, :]
+            else:
+                feat = stacked_demo_feat[(length-1):length, :]
+            if len(feat.size()) == 3:
+                feat = feat.unsqueeze(0)
+            
+            if self.pooling == 'max':
+                feat = torch.max(feat, 0)[0]
+            elif self.pooling == 'avg':
+                feat = torch.mean(feat, 0)
+            elif self.pooling == 'lstmavg':
+                lstm_out, hidden = self.lstm(feat.view(len(feat), 1, -1))
+                lstm_out = lstm_out.view(len(feat), -1)
+                feat = torch.mean(lstm_out, 0)
+            elif self.pooling == 'lstmlast':
+                lstm_out, hidden = self.lstm(feat.view(len(feat), 1, -1))
+                lstm_out = lstm_out.view(len(feat), -1)
+                feat = lstm_out[-1]
+            else:
+                raise ValueError
+
+
+            batch_demo_feat.append(feat)
+
+        demo_emb = torch.stack(batch_demo_feat, 0)
+        demo_emb = demo_emb.view(8,82, -1)
+        return demo_emb, batch_demo_feat
+

watch_and_help/watch_strategy_full/predicate-train-strategy.py

@@ -0,0 +1,368 @@
+import resource
+import time
+from termcolor import colored
+import torch
+from torch.utils.data import DataLoader
+from helper import Constant, LinearStep
+from predicate.utils import save, setup, save_checkpoint
+from predicate.utils import summary, write_prob, summary_eval, write_prob_strategy
+import random
+import json
+import pickle
+import numpy as np
+
+
+topk = 1
+p_th = 0.5
+
+
+def print_output(args, outputs, targets, file_names, test_dset):
+    goal_predicates = test_dset.goal_predicates
+    goal_predicates = {v:k for k,v in goal_predicates.items()}
+    json_output = {}
+
+    for i, target in enumerate(targets):      
+        if args.inference == 0:
+            p = random.uniform(0, 1)
+            if p>p_th:
+                continue
+
+        file_name = file_names[i]
+        output = outputs[i]
+
+        if args.multi_classifier:
+            output = torch.Tensor(output).view(-1, len(test_dset.goal_predicates), test_dset.max_subgoal_length+1)
+            target = torch.Tensor(target).view(-1, len(test_dset.goal_predicates))
+        else:
+            output = torch.Tensor(output).view(-1, test_dset.max_goal_length, len(test_dset.goal_predicates))
+            target = torch.Tensor(target).view(-1, test_dset.max_goal_length)
+
+        output = output.numpy()
+        target = target.numpy()
+
+        if args.inference == 0:
+            target_inference = [target[0]]
+            output_inference = [output[0]]
+            file_name_inference = [file_name[0]]
+        else:
+            target_inference = target
+            output_inference = output
+            file_name_inference = file_name
+
+
+        for (target_j, output_j, file_name_j) in zip(target_inference, output_inference, file_name_inference):
+            ## only show the fist sample in each minibatch
+            assert file_name_j not in json_output
+            json_output[file_name_j] = {}
+            json_output[file_name_j]['ground_truth'] = []
+            json_output[file_name_j]['prediction'] = []
+            json_output[file_name_j]['ground_truth_id'] = []
+            json_output[file_name_j]['prediction_id'] = []
+
+            print('----------------------------------------------------------------------------------')
+            if args.multi_classifier:
+                assert len(target_j) == len(goal_predicates) == len(output_j)
+                for k, target_k in enumerate(target_j):
+                    output_k = output_j[k]
+                    strtar = ('tar: %s   %d' % (goal_predicates[k], target_k)).ljust(50, ' ')
+                    strpre = '| gen: %s   %d' % (goal_predicates[k], output_k.argmax())
+                    print(strtar+strpre)
+
+                    json_output[file_name_j]['ground_truth_id'].append(int(target_k))
+                    json_output[file_name_j]['prediction_id'].append(output_k.argmax())
+                    json_output[file_name_j]['ground_truth'].append(goal_predicates[k])
+                    json_output[file_name_j]['prediction'].append(goal_predicates[k])
+            else:
+                for k, target_k in enumerate(target_j):
+                    output_k = output_j[k]
+
+                    strtar = ('tar: %s' % goal_predicates[int(target_k)]).ljust(50, ' ')
+                    strpre = '| gen: %s' % goal_predicates[output_k.argmax()]
+                    print(strtar+strpre)
+
+                    json_output[file_name_j]['ground_truth_id'].append(int(target_k))
+                    json_output[file_name_j]['prediction_id'].append(output_k.argmax())
+                    json_output[file_name_j]['ground_truth'].append(goal_predicates[int(target_k)])
+                    json_output[file_name_j]['prediction'].append(goal_predicates[output_k.argmax()])
+
+            print('----------------------------------------------------------------------------------')
+
+    if args.inference == 1:
+        if args.single:
+            pickle.dump( json_output, open( "dataset/test_output_"+args.resume.split('/')[-2]+"_single_task.p", "wb" ) )
+        else:
+            pickle.dump( json_output, open( "dataset/test_output_"+args.resume.split('/')[-2]+"_multiple_task.p", "wb" ) )
+
+def run_one_iteration(model, optim, batch_data, train_args, args):
+    model.train()
+    optim.zero_grad()
+    loss, info = model(batch_data, **train_args)
+    loss.backward()
+    optim.step()
+    return batch_data, info, loss
+
+
+def train(
+        args,
+        model,
+        optim,
+        train_loader,
+        test_loader,
+        val_loader,
+        checkpoint_dir,
+        writer,
+        train_dset,
+        test_dset,
+        task):
+    # Train
+    print(colored('Start training...', 'red'))
+    # loader for the testing set
+    def _loader():
+        while True:
+            for batch_data in test_loader:
+                yield batch_data
+
+    get_next_data_fn = _loader().__iter__().__next__
+    train_args = {}
+    
+    if args.inference == 1:
+        info = summary(
+            args,
+            writer,
+            None,
+            None,
+            model,
+            test_loader,
+            'test')      
+        print('test top1', info['top1'])
+        write_prob(info, args)               
+
+    def _train_loop(task):
+        iter = 0
+        summary_t1 = time.time()
+
+        test_best_top1 = 0
+        print('start while')
+        print('train iterations: ',args.train_iters)
+        while iter <= args.train_iters:
+            for batch_data in train_loader:                
+                results = run_one_iteration(model, optim, batch_data, train_args, args)
+                batch_data, info, loss = results
+
+                if iter % 10 == 0:
+                    print('%s: training %d / %d: loss %.4f: acc %.4f' % (args.checkpoint, iter, len(train_loader), loss, info['top1']))
+
+                    fps = 10. / (time.time() - summary_t1)
+                    info = summary(
+                        args,
+                        writer,
+                        info,
+                        train_args,
+                        model,
+                        None,
+                        'train',
+                        fps=fps)
+                    if iter > 0:
+                        summary_t1 = time.time()
+                
+                if iter % (len(train_loader)*1) == 0 and iter>0:
+                    info = summary(
+                        args,
+                        writer,
+                        None,
+                        None,
+                        model,
+                        test_loader,
+                        'test')
+
+                    if info['top1']>test_best_top1:
+                        test_best_top1 = info['top1']
+                        save(args, iter, checkpoint_dir, model, task)
+                    save_checkpoint(args, iter, checkpoint_dir, model, task)
+
+                iter += 1
+    print('start train loop')
+    _train_loop(task)
+    print('train loop done')
+
+
+def main():    
+    args, checkpoint_dir, writer, model_config = setup(train=True)    
+    print(args)
+
+    from predicate.demo_dataset_graph_strategy_test import get_dataset
+    from predicate.demo_dataset_graph_strategy_test import collate_fn
+    from predicate.demo_dataset_graph_strategy_test import to_cuda_fn
+    
+    #strategy inference        
+    if args.inference == 2: # 0: not infer, 1: infer, 2: strategy infer   
+        from network.encoder_decoder import ActionDemo2Predicate     
+        test_tasks = ['put_fridge', 'put_dishwasher', 'read_book', 'prepare_food', 'setup_table']
+        new_test_tasks = ['put_fridge', 'put_dishwasher', 'read_book']
+        
+        train_dsets = []
+        test_dsets = []
+        new_test_dsets = []
+        models = []
+        train_loaders = []
+        test_loaders = []
+        val_loaders = []
+        for i in range(len(new_test_tasks)):
+            loss_weights = np.load('dataset/watch_data/loss_weight_'+test_tasks[i]+'_new_test_task'+'.npy')
+            train_dset, test_dset, new_test_dset = get_dataset(args, new_test_tasks[i], train=True )
+            train_dsets.append(train_dset)
+            test_dsets.append(test_dset)
+            new_test_dsets.append(new_test_dset)
+            model = ActionDemo2Predicate(args, train_dset, loss_weights, **model_config)          
+            model.load(args.checkpoint+'/demo2predicate-checkpoint_model_'+new_test_tasks[i]+'.ckpt', True)
+            model.cuda()
+            model.eval()
+            models.append(model)
+            
+            train_loader = DataLoader(
+                dataset=train_dset,
+                batch_size=args.batch_size,
+                shuffle=True,
+                num_workers=args.n_workers,
+                drop_last=True)
+            if args.testset == 'test_task':
+                test_loader = DataLoader(
+                    dataset=test_dset,
+                    batch_size=args.batch_size,
+                    shuffle=False,
+                    num_workers=0,
+                    drop_last=True)
+
+                val_loader = DataLoader(
+                    dataset=test_dset,
+                    batch_size=args.batch_size,
+                    shuffle=False,
+                    num_workers=0,
+                    drop_last=True)
+            if args.testset == 'new_test_task':
+                test_loader = DataLoader(
+                    dataset=new_test_dset,
+                    batch_size=args.batch_size,
+                    shuffle=False,
+                    num_workers=0,
+                    drop_last=True)
+
+                val_loader = DataLoader(
+                    dataset=new_test_dset,
+                    batch_size=args.batch_size,
+                    shuffle=False,
+                    num_workers=0,
+                    drop_last=True)
+                            
+            train_loaders.append(train_loader)
+            test_loaders.append(test_loader)
+            val_loaders.append(val_loader)
+        
+        
+        for i in range(len(models)):
+            infos = []
+            for j in range(len(test_loaders)):
+                info = summary_eval(
+                    models[i],
+                    test_loaders[j],
+                    test_loaders[j].dataset)
+                print('test top1', info['top1'])
+                infos.append(info)
+            total_info = {
+                "prob": np.concatenate((infos[0]["prob"], infos[1]["prob"], infos[2]["prob"]), axis=0),               
+                "target": np.concatenate((infos[0]["target"], infos[1]["target"], infos[2]["target"]), axis=0), #batch_target.cpu().numpy(),
+                "task_name": np.concatenate((infos[0]["task_name"], infos[1]["task_name"], infos[2]["task_name"]), axis=0), #batch_task_name,
+                "action_id": np.concatenate((infos[0]["action_id"], infos[1]["action_id"], infos[2]["action_id"]), axis=0) #batch_action_id.cpu().numpy()
+            }           
+            write_prob_strategy(total_info, test_tasks[i], args)
+                    
+    else:           
+        print('get dataset')
+        test_tasks = ['put_dishwasher', 'read_book', 'put_fridge', 'prepare_food', 'setup_table']
+        new_test_tasks = ['put_dishwasher', 'read_book', 'put_fridge']
+        for i in range(len(new_test_tasks)):
+            train_dset, test_dset, new_test_dset = get_dataset(args, test_tasks[i], train=True )
+            print('train set len:',len(train_dset))            
+
+            train_loader = DataLoader(
+                dataset=train_dset,
+                batch_size=args.batch_size,
+                shuffle=True,
+                num_workers=args.n_workers,
+                drop_last=True)
+
+            if args.single:
+                test_loader = DataLoader(
+                    dataset=new_test_dset,
+                    batch_size=args.batch_size,
+                    shuffle=True,
+                    num_workers=0,
+                    drop_last=True)
+
+                val_loader = DataLoader(
+                    dataset=test_dset, 
+                    batch_size=args.batch_size,
+                    shuffle=True,
+                    num_workers=0,
+                    drop_last=True)
+            else:
+                test_loader = DataLoader(
+                    dataset=new_test_dset,
+                    batch_size=args.batch_size,
+                    shuffle=True,
+                    num_workers=0,
+                    drop_last=True)
+
+                val_loader = DataLoader(
+                    dataset=test_dset,
+                    batch_size=args.batch_size,
+                    shuffle=True,
+                    num_workers=0,
+                    drop_last=True)
+
+            # initialize model
+            loss_weights = np.load('dataset/watch_data/loss_weight_'+test_tasks[i]+'_train_task'+'.npy')
+            if args.inputtype=='graphinput':
+                from network.encoder_decoder import GraphDemo2Predicate
+                model = GraphDemo2Predicate(args, train_dset, **model_config)
+            elif args.inputtype=='actioninput':
+                from network.encoder_decoder import ActionDemo2Predicate
+                model = ActionDemo2Predicate(args, train_dset, loss_weights, **model_config)
+
+            if args.resume!='':
+                model.load(args.resume, True)
+
+
+            optim = torch.optim.Adam(
+                filter(
+                    lambda p: p.requires_grad,
+                    model.parameters()),
+                args.model_lr_rate)
+            if args.gpu_id is not None:
+                model.cuda()
+            
+            # main loop
+            train(
+                args,
+                model,
+                optim,
+                train_loader,
+                test_loader,
+                val_loader,
+                checkpoint_dir,
+                writer,
+                train_dset, 
+                test_dset,
+                test_tasks[i])           
+
+rlimit = resource.getrlimit(resource.RLIMIT_NOFILE)
+resource.setrlimit(resource.RLIMIT_NOFILE, (1024 * 4, rlimit[1]))
+
+
+if __name__ == '__main__':
+    from multiprocessing import set_start_method
+    try:
+        set_start_method('spawn')
+    except RuntimeError:
+        pass
+    main()

watch_and_help/watch_strategy_full/predicate/demo_dataset_graph_strategy_test.py

@@ -0,0 +1,306 @@
+import os
+import random
+import copy
+import json
+
+
+import numpy as np
+from termcolor import colored
+from glob import glob
+import pickle
+
+
+import torch
+import torch.nn.functional as F
+from torch.utils.data.dataset import Dataset
+
+
+
+################################
+# Demonstration
+################################
+
+def get_dataset(args, task_name, train):   
+    train_data, test_data, new_test_data, train_action_gt, test_action_gt, new_test_action_gt, train_task_name, test_task_name, new_test_task_name, train_action_id, test_action_id, new_test_action_id, action_predicates, all_action, all_object, goal_objects, goal_targets, goal_predicates, max_goal_length, max_action_length, max_node_length, max_subgoal_length = gather_data(args, task_name)    
+    train_dset = demo_dset(args, train_data, train_action_gt, train_task_name, train_action_id, action_predicates, all_action, all_object, goal_objects, goal_targets, goal_predicates, max_goal_length, max_action_length, max_node_length, max_subgoal_length)
+    test_dset = demo_dset(args, test_data, test_action_gt, test_task_name, test_action_id, action_predicates, all_action, all_object, goal_objects, goal_targets, goal_predicates, max_goal_length, max_action_length, max_node_length, max_subgoal_length)
+    new_test_dset = demo_dset(args, new_test_data, new_test_action_gt, new_test_task_name, new_test_action_id, action_predicates, all_action, all_object, goal_objects, goal_targets, goal_predicates, max_goal_length, max_action_length, max_node_length, max_subgoal_length)
+    return train_dset, test_dset, new_test_dset
+
+
+def collate_fn(data_list):
+    graph_data = [data[0] for data in data_list]
+    batch_goal_index = [data[1] for data in data_list]
+    batch_valid_action_with_walk_index = [data[2] for data in data_list]
+
+    if len(graph_data[0])==3:
+        batch_graph_length = [d[0] for d in graph_data]
+        batch_graph_input = [d[1] for d in graph_data]
+        batch_file_name = [d[2] for d in graph_data]
+    else:
+        batch_graph_length = [d[0] for d in graph_data]
+        batch_file_name = [d[1] for d in graph_data]
+
+    if len(graph_data[0])==3:
+        batch_demo_data = (
+            np.arange(len(batch_graph_length)), 
+            batch_graph_length, 
+            batch_graph_input,
+            batch_file_name
+        )
+    else:
+        batch_demo_data = (
+            np.arange(len(batch_graph_length)), 
+            batch_graph_length,
+            batch_file_name
+        )
+
+    return batch_demo_data, batch_goal_index, batch_valid_action_with_walk_index
+
+
+def to_cuda_fn(data):
+    batch_demo_data, batch_goal_index, batch_valid_action_with_walk_index = data
+
+    if len(batch_demo_data)==4:
+        batch_demo_index, batch_graph_length, batch_graph_input, batch_file_name = batch_demo_data
+
+        batch_graph_input_class_objects = [[torch.tensor(j['class_objects']).cuda() for j in i] for i in batch_graph_input]
+        batch_graph_input_object_coords = [[torch.tensor(j['object_coords']).cuda() for j in i] for i in batch_graph_input]
+        batch_graph_input_states_objects = [[torch.tensor(j['states_objects']).cuda() for j in i] for i in batch_graph_input]
+        batch_graph_input_mask_object = [[torch.tensor(j['mask_object']).cuda() for j in i] for i in batch_graph_input]
+        
+        batch_graph_input = { 'class_objects': batch_graph_input_class_objects, 
+                                'object_coords': batch_graph_input_object_coords,
+                                'states_objects': batch_graph_input_states_objects,
+                                'mask_object': batch_graph_input_mask_object}
+
+    else:
+        batch_demo_index, batch_graph_length, batch_file_name = batch_demo_data
+    
+
+    batch_goal_index = [torch.tensor(i).cuda().long() for i in batch_goal_index]
+    batch_valid_action_with_walk_index = [torch.tensor(i).cuda().long() for i in batch_valid_action_with_walk_index]
+    
+    if len(batch_demo_data)==4:
+        batch_demo_data = (
+            batch_demo_index,
+            batch_graph_length, 
+            batch_graph_input,
+            batch_file_name
+        )
+    else:
+        batch_demo_data = (
+            batch_demo_index,
+            batch_graph_length,
+            batch_file_name
+        )
+
+    return batch_demo_data, batch_goal_index, batch_valid_action_with_walk_index
+
+def one_hot(states, graph_node_states):
+    one_hot = np.zeros(len(graph_node_states))
+    for state in states:
+        one_hot[graph_node_states[state]] = 1
+    return one_hot
+
+def gather_data(args, task):    
+    meta_data_path = 'dataset/watch_data/metadata.json'
+    data_path_new_test = 'dataset/watch_data/action/new_test_task_' + task + '_strategy.json'
+    data_path_test = 'dataset/watch_data/action/test_task_' + task + '_strategy.json'
+    data_path_train = 'dataset/watch_data/action/train_task_' + task + '_strategy.json'
+
+    
+    with open(data_path_new_test, 'r') as f:
+        new_test_data = json.load(f)
+        
+    with open(data_path_test, 'r') as f:
+        test_data = json.load(f)
+    
+    #if os.path.exists(data_path):
+        #print('load gather_data, this may take a while...', data_path)
+    with open(data_path_train, 'r') as f:
+        data = json.load(f)
+      
+    # temporarily set data_path to test data to test training
+    train_data = data
+
+    with open(meta_data_path, 'r') as f:
+        data = json.load(f)
+            
+        action_predicates = data['action_predicates']
+        all_action = data['all_action']
+        all_object = data['all_object']
+        
+        goal_objects = data['goal_objects']
+        goal_targets = data['goal_targets']
+        goal_predicates = data['goal_predicates']
+            
+        graph_class_names = data['graph_class_names']
+        graph_node_states = data['graph_node_states']
+
+        max_goal_length = data['max_goal_length']
+        max_action_length = data['max_action_length']
+        max_node_length = data['max_node_length']
+            
+
+    ## -----------------------------------------------------------------------------
+    ## add action, goal, and graph node index
+    ## -----------------------------------------------------------------------------
+    max_subgoal_length = 1
+    
+    train_task_name = np.array(train_data['task_name'])
+    
+    test_task_name = np.array(test_data['task_name'])
+    
+    new_test_task_name = np.array(new_test_data['task_name'])
+    
+    
+    for traintest in [train_data, test_data, new_test_data]:
+        for data in traintest['goal']:
+            ## goal
+            goal_index = []
+            subgoal_dict = {}
+            for subgoal in data:
+                goal_index.append(goal_predicates[subgoal])
+
+                if goal_predicates[subgoal] not in subgoal_dict:
+                    subgoal_dict[goal_predicates[subgoal]] = 1
+                else:
+                    subgoal_dict[goal_predicates[subgoal]] += 1
+
+            this_max_subgoal_length = np.max(list(subgoal_dict.values()))
+            if this_max_subgoal_length>max_subgoal_length:
+                max_subgoal_length = this_max_subgoal_length
+            
+
+            goal_index.sort()
+            if len(goal_index) < max_goal_length:
+                for i in range(max_goal_length-len(goal_index)):
+                    goal_index.append(0)                  
+
+        ## action gt       
+        for i in range(len(traintest['action_gt'])): # len(traintest['action_gt'])
+            action_name = traintest['action_gt'][i][0].split(' ')[0]
+            object_name = traintest['action_gt'][i][0].split(' ')[1]
+            predicate_name = ' '.join([action_name, object_name])
+            traintest['action_gt'][i] = action_predicates[predicate_name]
+                
+        ## action        
+        valid_action_with_walk_index = []
+        for i in range(len(traintest['valid_action_with_walks'])):   
+            actions_index = []            
+            for actions in traintest['valid_action_with_walks'][i]:
+                if actions!='None':
+                    action_name = actions[0].split(' ')[0]
+                    object_name = actions[0].split(' ')[1]
+                    predicate_name = ' '.join([action_name, object_name])
+                else:
+                    predicate_name = actions
+                actions_index.append(action_predicates[predicate_name])
+                
+            traintest['valid_action_with_walks'][i] = actions_index
+                      
+    print(len(train_data['action_gt']),np.array(train_data['action_gt']), type(train_data['action_gt']))
+    
+    
+    train_action_gt =  np.array(train_data['action_gt'])    
+    test_action_gt = np.array(test_data['action_gt'])    
+    new_test_action_gt = np.array(new_test_data['action_gt'])
+    
+    train_action_id =  np.array(train_data['action_id'])    
+    test_action_id = np.array(test_data['action_id'])    
+    new_test_action_id = np.array(new_test_data['action_id'])
+    
+    train_data = np.array(train_data['valid_action_with_walks'])        
+    test_data = np.array(test_data['valid_action_with_walks'])        
+    new_test_data = np.array(new_test_data['valid_action_with_walks'])
+    
+    
+    print('--------------------------------------------------------------------------------')
+    print('train_data', len(train_data), train_data.shape)
+    print('test_data', len(test_data), train_data.shape)
+    print('new_test_data', len(new_test_data), train_data.shape)
+    print('--------------------------------------------------------------------------------')
+    print('train_gt', len(train_action_gt), train_action_gt.shape)
+    print('test_gt', len(test_action_gt), test_action_gt.shape)
+    print('new_test_gt', len(new_test_action_gt), new_test_action_gt.shape)
+    print('--------------------------------------------------------------------------------')
+    print('train_task_name', len(train_task_name), train_task_name.shape)
+    print('test_task_name', len(test_task_name), test_task_name.shape)
+    print('new_test_task_name', len(new_test_task_name), new_test_task_name.shape)
+    print('--------------------------------------------------------------------------------')
+        
+    return train_data, test_data, new_test_data, train_action_gt, test_action_gt, new_test_action_gt, train_task_name, test_task_name, new_test_task_name, train_action_id, test_action_id, new_test_action_id, action_predicates, all_action, all_object, goal_objects, goal_targets, goal_predicates, max_goal_length, max_action_length, max_node_length, max_subgoal_length
+
+class demo_dset(Dataset):
+
+    def __init__(
+            self,
+            args,
+            data, 
+            gt,
+            task_name,
+            action_id,
+            #action_predicates, all_action, all_object, goal_objects, goal_targets, goal_predicates, graph_class_names, graph_node_states, max_goal_length, max_action_length, max_node_length, max_subgoal_length):
+            action_predicates, all_action, all_object, goal_objects, goal_targets, goal_predicates, max_goal_length, max_action_length, max_node_length, max_subgoal_length):
+
+
+        self.inputtype = args.inputtype
+        self.multi_classifier = args.multi_classifier
+        self.data = data
+        self.gt = gt
+        self.task_name = task_name
+        self.action_id = action_id
+        
+        self.max_action_len = 82
+        self.action_predicates = action_predicates
+        self.all_action = all_action
+        self.all_object = all_object
+        
+        self.goal_objects = goal_objects
+        self.goal_targets = goal_targets
+        self.goal_predicates = goal_predicates
+        self.num_goal_predicates = len(goal_predicates)
+
+        self.max_goal_length = max_goal_length
+        self.max_action_length = max_action_length
+        self.max_subgoal_length = max_subgoal_length
+
+        if self.inputtype=='graphinput':
+            self.graph_class_names = graph_class_names
+            self.graph_node_states = graph_node_states
+            self.num_node_states = len(graph_node_states)
+            self.max_node_length = max_node_length
+
+
+        print('-----------------------------------------------------------------------------')
+        print('num_goal_predicates', self.num_goal_predicates)
+        print('max_goal_length', self.max_goal_length)
+        print('max_action_length', max_action_length)
+        
+        if self.inputtype=='graphinput':
+            print('num_node_states', self.num_node_states)
+            print('max_node_length', max_node_length)
+        print('-----------------------------------------------------------------------------')
+
+        
+    def __getitem__(self, index):
+        data = self.data[index]
+        gt = self.gt[index]
+        task_name = self.task_name[index]
+        action_id = self.action_id[index]
+        
+        return data, gt, task_name, action_id
+
+    def __len__(self):
+        return len(self.data)
+
+    def _preprocess_one_data(self, data):       
+        action_gt = data['action_gt']
+        valid_action_with_walk_index = data['valid_action_with_walks']
+        action_length = len(valid_action_with_walk_index)
+        inputdata = (action_length, 'actions')
+        data = [inputdata, action_gt, valid_action_with_walk_index]
+        return data
+    
+        

watch_and_help/watch_strategy_full/predicate/utils.py

@@ -0,0 +1,297 @@
+import argparse
+import random
+import time
+import os
+import json
+import numpy as np
+
+
+import torch
+from torch.utils.tensorboard import SummaryWriter
+from helper import to_cpu, average_over_list, writer_helper
+import csv
+import pathlib
+
+def grab_args():
+
+    def str2bool(v):
+        return v.lower() == 'true'
+
+    parser = argparse.ArgumentParser(description='')
+    parser.add_argument('--seed', type=int, default=123, help='random seed')
+    parser.add_argument('--verbose', type=str2bool, default=False)
+    parser.add_argument('--debug', type=str2bool, default=False)
+    parser.add_argument('--prefix', type=str, default='test')
+    parser.add_argument('--checkpoint', type=str, default=None)
+    parser.add_argument('--n_workers', type=int, default=0)
+    parser.add_argument('--train_iters', type=int, default=2e4)
+    parser.add_argument('--inputtype', type=str, default='actioninput')
+    parser.add_argument('--resume', type=str, default='')
+    parser.add_argument('--dropout', type=float, default=0)
+    parser.add_argument('--inference', type=int, default=0)
+    parser.add_argument('--single', type=int, default=0)
+    parser.add_argument('--loss_type', type=str, default='regu') #regu or ce
+    parser.add_argument('--testset', type=str, default='test') # test: test set 1, new_test: test set 2
+
+    # model config
+    parser.add_argument(
+        '--model_type',
+        type=str,
+        default='max')
+    parser.add_argument('--embedding_dim', type=int, default=100)
+    parser.add_argument('--predicate_hidden', type=int, default=128)
+    parser.add_argument('--demo_hidden', type=int, default=128)
+    parser.add_argument('--multi_classifier', type=int, default=0)
+    parser.add_argument('--transformer_nhead', type=int, default=2)
+    
+    # train config
+    parser.add_argument(
+        '--gpu_id',
+        metavar='N',
+        type=str,
+        nargs='+',
+        help='specify the gpu id')
+    parser.add_argument('--batch_size', type=int, default=2)
+    parser.add_argument('--model_lr_rate', type=float, default=3e-4)
+
+    args = parser.parse_args()
+    return args
+
+
+def setup(train):
+    def _basic_setting(args):
+
+        # set seed
+        torch.manual_seed(args.seed)
+        random.seed(args.seed)
+        np.random.seed(args.seed)
+
+        if args.gpu_id is None:
+            os.environ['CUDA_VISIBLE_DEVICES'] = ''
+            args.__dict__.update({'cuda': False})
+        else:
+            os.environ['CUDA_VISIBLE_DEVICES'] = ', '.join(args.gpu_id)
+            args.__dict__.update({'cuda': True})
+            torch.cuda.manual_seed_all(args.seed)
+
+        if args.debug:
+            args.verbose = True
+
+    def _basic_checking(args):
+        pass
+
+    def _create_checkpoint_dir(args):
+         # setup checkpoint_dir
+        if args.debug:
+            checkpoint_dir = 'debug'
+        elif train:
+            checkpoint_dir = 'checkpoint_dir'
+        else:
+            checkpoint_dir = 'testing_dir'
+
+        checkpoint_dir = os.path.join(checkpoint_dir, 'demo2predicate')
+
+        args_dict = args.__dict__
+        keys = sorted(args_dict)
+        prefix = ['{}-{}'.format(k, args_dict[k]) for k in keys]
+        prefix.remove('debug-{}'.format(args.debug))
+        prefix.remove('checkpoint-{}'.format(args.checkpoint))
+        prefix.remove('gpu_id-{}'.format(args.gpu_id))
+
+        checkpoint_dir = os.path.join(checkpoint_dir, *prefix)
+
+        checkpoint_dir += '/{}'.format(time.strftime("%Y%m%d-%H%M%S"))
+
+        return checkpoint_dir
+
+    def _make_dirs(checkpoint_dir, tfboard_dir):
+        if not os.path.exists(checkpoint_dir):
+            os.makedirs(checkpoint_dir)
+        if not os.path.exists(tfboard_dir):
+            os.makedirs(tfboard_dir)
+
+    def _print_args(args):
+        args_str = ''
+        with open(os.path.join(checkpoint_dir, 'args.txt'), 'w') as f:
+            for k, v in args.__dict__.items():
+                s = '{}: {}'.format(k, v)
+                args_str += '{}\n'.format(s)
+                print(s)
+                f.write(s + '\n')
+        print("All the data will be saved in", checkpoint_dir)
+        return args_str
+
+    args = grab_args()
+    _basic_setting(args)
+    _basic_checking(args)
+
+    checkpoint_dir = args.checkpoint
+    tfboard_dir = os.path.join(checkpoint_dir, 'tfboard')
+
+    _make_dirs(checkpoint_dir, tfboard_dir)
+    args_str = _print_args(args)
+
+    writer = SummaryWriter(tfboard_dir)
+    writer.add_text('args', args_str, 0)
+    writer = writer_helper(writer)
+
+    model_config = {
+        "model_type": args.model_type,
+        "embedding_dim": args.embedding_dim,
+        "predicate_hidden": args.predicate_hidden,
+    }
+    model_config.update({"demo_hidden": args.demo_hidden})
+
+    return args, checkpoint_dir, writer, model_config
+
+
+def summary(
+        args,
+        writer,
+        info,
+        train_args,
+        model,
+        test_loader,
+        postfix, 
+        fps=None):
+    
+    if postfix == 'train':
+        model.write_summary(writer, info, postfix=postfix)
+    elif postfix == 'val':
+        info = summary_eval(
+            model,
+            test_loader,
+            test_loader.dataset)
+        model.write_summary(writer, info, postfix=postfix)
+    elif postfix == 'test':
+        info = summary_eval(
+            model,
+            test_loader,
+            test_loader.dataset)
+
+        model.write_summary(writer, info, postfix=postfix)
+    else:
+        raise ValueError
+
+    if fps:
+        writer.scalar_summary('General/fps', fps)
+
+    return info
+
+def summary_eval(
+        model,
+        loader,
+        dset):
+
+    model.eval()
+    print(len(loader))
+    with torch.no_grad():
+
+        loss_list = []
+        top1_list = []
+        iter = 0
+        action_id_list = []
+        
+        prob = []
+        target = []
+        file_name = []
+        for batch_data in loader:
+            loss, info = model(batch_data)
+            loss_list.append(loss.cpu().item())
+            top1_list.append(info['top1'])
+
+            prob.append(info['prob'])
+            target.append(info['target'])
+            file_name.append(info['task_name'])
+            action_id_list.append(info['action_id'])
+
+            if iter%10==0:
+                print('testing %d / %d: loss %.4f: acc %.4f' % (iter, len(loader), loss, info['top1']))
+
+            iter += 1
+
+    info = {"loss": sum(loss_list)/ len(loss_list), "top1": sum(top1_list)/ len(top1_list), "prob": prob, "target": target, "task_name": file_name, "action_id": action_id_list}
+    return info
+
+def write_prob(info, args):
+    temp_prob_list = []
+    temp_action_id_list = []
+    temp_task_name_list = []
+    temp_target_list = []
+    for i in range(len(info['prob'])):
+        for j in range(len(info['prob'][i])):
+            temp_prob_list.append(info['prob'][i][j])
+            
+    for i in range(len(info['action_id'])):
+        for j in range(len(info['action_id'][i])):        
+            temp_action_id_list.append(info['action_id'][i][j])
+            
+    for i in range(len(info['task_name'])):
+        for j in range(len(info['task_name'][i])):        
+            temp_task_name_list.append(info['task_name'][i][j])
+            
+    for i in range(len(info['target'])):
+        for j in range(len(info['target'][i])):        
+            temp_target_list.append(info['target'][i][j])
+            
+    prob = np.array(temp_prob_list)
+    action_id = np.array(temp_action_id_list)
+    task_name = np.array(temp_task_name_list)
+    target = np.array(temp_target_list)    
+
+    write_data = np.concatenate((np.reshape(action_id, (-1, 1)), prob, np.reshape(task_name, (-1, 1)), np.reshape(target, (-1, 1))), axis=1)
+    import pandas as pd 
+    head = []
+    for j in range(79):
+        head.append('act'+str(j+1))
+    head.append('task_name')
+    head.append('gt')
+    head.insert(0,'action_id')
+    pd.DataFrame(write_data).to_csv("prediction/" + args.model_type + "/" + task_name[0] + "_full.csv", header=head)
+
+def write_prob_strategy(info, model_name, args):
+    temp_prob_list = []
+    temp_action_id_list = []
+    temp_task_name_list = []
+    temp_target_list = []
+    for i in range(len(info['prob'])):
+        for j in range(len(info['prob'][i])):
+            temp_prob_list.append(info['prob'][i][j])
+            
+    for i in range(len(info['action_id'])):
+        for j in range(len(info['action_id'][i])):        
+            temp_action_id_list.append(info['action_id'][i][j])
+            
+    for i in range(len(info['task_name'])):
+        for j in range(len(info['task_name'][i])):        
+            temp_task_name_list.append(info['task_name'][i][j])
+            
+    for i in range(len(info['target'])):
+        for j in range(len(info['target'][i])):        
+            temp_target_list.append(info['target'][i][j])
+            
+    prob = np.array(temp_prob_list)
+    action_id = np.array(temp_action_id_list)
+    task_name = np.array(temp_task_name_list)
+    target = np.array(temp_target_list)
+    
+
+    write_data = np.concatenate((np.reshape(action_id, (-1, 1)), prob, np.reshape(task_name, (-1, 1)), np.reshape(target, (-1, 1))), axis=1)
+    import pandas as pd 
+    head = []
+    for j in range(79):
+        head.append('act'+str(j+1))
+    head.append('task_name')
+    head.append('gt')
+    head.insert(0,'action_id')
+    path = pathlib.Path("stan/prediction/" + args.testset + "/" +  args.model_type)
+    path.mkdir(parents=True, exist_ok=True)
+    pd.DataFrame(write_data).to_csv("stan/prediction/" + args.testset + "/" +  args.model_type + "/model_" + model_name + '_strategy_' + task_name[0] + ".csv", header=head)
+
+def save(args, i, checkpoint_dir, model, task):    
+    save_path = '{}/demo2predicate-{}{}.ckpt'.format(checkpoint_dir, 'best_model_', task)
+    model.save(save_path, True)
+
+def save_checkpoint(args, i, checkpoint_dir, model, task):    
+    save_path = '{}/demo2predicate-{}{}.ckpt'.format(checkpoint_dir, 'checkpoint_model_', task)
+    model.save(save_path, True)