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)