train_loader = DataLoader(train_data,
args.batch_sz,
shuffle=True,
collate_fn=get_collate_fn(args.cuda))
valid_data = TalkTheWalkLanguage(data_dir, 'valid')
valid_loader = DataLoader(valid_data,
args.batch_sz,
collate_fn=get_collate_fn(args.cuda))
tourist = TouristLanguage(args.act_emb_sz,
args.act_hid_sz,
len(train_data.act_dict),
args.obs_emb_sz,
args.obs_hid_sz,
len(train_data.map.landmark_dict),
args.decoder_emb_sz,
args.decoder_hid_sz,
len(train_data.dict),
start_token=train_data.dict.tok2i[START_TOKEN],
end_token=train_data.dict.tok2i[END_TOKEN])
opt = optim.Adam(tourist.parameters())
if args.cuda:
tourist = tourist.cuda()
best_val = 1e10
for epoch in range(1, args.num_epochs):
train_loss = eval_epoch(train_loader, tourist, opt=opt)
train_data.dict = dictionary
valid_data = TalkTheWalkEmergent(data_dir, 'valid', T=args.T)
valid_data.dict = dictionary
test_data = TalkTheWalkEmergent(data_dir, 'test', T=args.T)
test_data.dict = dictionary
elif args.trajectories == 'human':
train_data = TalkTheWalkLanguage(data_dir, 'train')
valid_data = TalkTheWalkLanguage(data_dir, 'valid')
test_data = TalkTheWalkLanguage(data_dir, 'test')
train_loader = DataLoader(train_data, args.batch_sz, collate_fn=get_collate_fn(args.cuda))
valid_loader = DataLoader(valid_data, args.batch_sz, collate_fn=get_collate_fn(args.cuda))
test_loader = DataLoader(test_data, args.batch_sz, collate_fn=get_collate_fn(args.cuda))
tourist = TouristLanguage.load(args.tourist_model)
if args.guide_model is not None:
guide = GuideLanguage.load(args.guide_model)
else:
guide = GuideLanguage(128, 256, len(train_data.dict), apply_masc=True, T=3)
if args.cuda:
tourist = tourist.cuda()
guide = guide.cuda()
if args.train_guide:
logger.info('Train guide (supervised)')
g_opt = optim.Adam(guide.parameters())
if args.train_tourist:
args.decoding_strategy = 'greedy'
from ttw.models import TouristLanguage
parser = argparse.ArgumentParser()
parser.add_argument('--data-dir',
type=str,
default='./data',
help='Path to talkthewalk dataset')
parser.add_argument('--tourist-model',
type=str,
default='Path to tourist checkpoint')
args = parser.parse_args()
train_data = TalkTheWalkLanguage(args.data_dir, 'train')
tourist_sl = TouristLanguage.load(args.tourist_model).cuda()
indices = []
for _ in range(5):
indices.append(random.randint(0, len(train_data) - 1))
indices = [32, 245, 560, 750, 1200, 2467, 2500]
print('supervised, greedy')
tourist_sl.show_samples(train_data,
indices=indices,
decoding_strategy='greedy')
print()
print()
print('supervised, beam')
tourist_sl.show_samples(train_data,
indices=indices,