369 lines
14 KiB
Python
369 lines
14 KiB
Python
|
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()
|