def main(): args = get_args() src_vocab_size = len( pickle.load(open(args.data_bin + '/dict' + "." + args.src_lang, 'rb')).keys()) tgt_vocab_size = len( pickle.load(open(args.data_bin + '/dict' + '.' + args.tgt_lang, 'rb')).keys()) device = 'cuda' model = Transformer(src_vocab_size=src_vocab_size, tgt_vocab_size=tgt_vocab_size, encoder_layer_num=args.encoder_layer_num, decoder_layer_num=args.decoder_layer_num, hidden_size=args.hidden_size, feedback_size=args.feedback, num_head=args.num_head, dropout=args.dropout, device=device) optim = Optim(Adam(model.parameters(), betas=(0.9, 0.98), eps=1e-9), warmup_step=4000, d_model=args.hidden_size) train_loader = DataLoader(Dataload(args.data_bin + '/' + 'train', args.src_lang, args.tgt_lang), batch_size=args.batch_size, collate_fn=collate_fn, shuffle=True) # optim = Adam(model.parameters(), lr=5e-6) test_loader = DataLoader(Dataload(args.data_bin + '/' + 'test', args.src_lang, args.tgt_lang), batch_size=args.batch_size, collate_fn=collate_fn) valid_loader = DataLoader(Dataload(args.data_bin + '/' + 'valid', args.src_lang, args.tgt_lang), batch_size=args.batch_size, collate_fn=collate_fn) best_loss = 1e4 model = model.to(device) # model.load_state_dict(torch.load('best_model.pkl')) for i in range(args.epoch): train(i, model, data_loader=train_loader, optim=optim, device=device) with torch.no_grad(): best_loss = eval(i, model, valid_loader, best_loss, device)
args = parser.parse_args() src_dict = pickle.load( open(add_(args.bin_path) + 'dict.' + args.src_lang, 'rb')) trg_dict = pickle.load( open(add_(args.bin_path) + 'dict.' + args.tgt_lang, 'rb')) model = Transformer(src_vocab_size=len(src_dict.keys()), tgt_vocab_size=len(trg_dict.keys()), encoder_layer_num=6, decoder_layer_num=6, hidden_size=512, feedback_size=2048, num_head=8, dropout=0.1, device=device) model = model.to(device) model.load_state_dict(torch.load(args.model_path)) dataload = DataLoader(Dataload(add_(args.bin_path) + 'test', src=args.src_lang, trg=args.tgt_lang), batch_size=32, collate_fn=collate_fn) real = [] predict = [] pbtr = tqdm(total=len(dataload)) with torch.no_grad(): model.eval() for src, trg in dataload: src = src.to(device) predicts = beamsearch(model, src, 1, 100, device=device)