def train_entry(config):
from models import QANet
with open(config.word_emb_file, "r") as fh:
word_mat = np.array(json.load(fh), dtype=np.float32)
with open(config.char_emb_file, "r") as fh:
char_mat = np.array(json.load(fh), dtype=np.float32)
with open(config.train_eval_file, "r") as fh:
train_eval_file = json.load(fh)
with open(config.dev_eval_file, "r") as fh:
dev_eval_file = json.load(fh)
print("Building model...")
train_dataset = SQuADDataset(config.train_record_file, config.num_steps,
config.batch_size)
dev_dataset = SQuADDataset(config.dev_record_file, config.test_num_batches,
config.batch_size)
lr = config.learning_rate
base_lr = 1.0
warm_up = config.lr_warm_up_num
model = QANet(word_mat, char_mat).to(device)
ema = EMA(config.ema_decay)
for name, p in model.named_parameters():
if p.requires_grad: ema.set(name, p)
params = filter(lambda param: param.requires_grad, model.parameters())
optimizer = optim.Adam(lr=base_lr,
betas=(config.beta1, config.beta2),
eps=1e-7,
weight_decay=3e-7,
params=params)
cr = lr / log2(warm_up)
scheduler = optim.lr_scheduler.LambdaLR(
optimizer,
lr_lambda=lambda ee: cr * log2(ee + 1) if ee < warm_up else lr)
L = config.checkpoint
N = config.num_steps
best_f1 = best_em = patience = 0
for iter in range(0, N, L):
train(model, optimizer, scheduler, ema, train_dataset, iter, L)
valid(model, train_dataset, train_eval_file)
metrics = test(model, dev_dataset, dev_eval_file)
print("Learning rate: {}".format(scheduler.get_lr()))
dev_f1 = metrics["f1"]
dev_em = metrics["exact_match"]
if dev_f1 < best_f1 and dev_em < best_em:
patience += 1
if patience > config.early_stop: break
else:
patience = 0
best_f1 = max(best_f1, dev_f1)
best_em = max(best_em, dev_em)
fn = os.path.join(config.save_dir,
f"model_iter={iter}_best_f1={best_f1}.pt")
torch.save(model, fn)
def train_entry(config):
from models import QANet
with open(config.word_emb_file, "rb") as fh:
word_mat = np.array(json.load(fh), dtype=np.float32)
with open(config.char_emb_file, "rb") as fh:
char_mat = np.array(json.load(fh), dtype=np.float32)
with open(config.dev_eval_file, "r") as fh:
dev_eval_file = json.load(fh)
print("Building model...")
train_dataset = get_loader(config.train_record_file, config.batch_size)
dev_dataset = get_loader(config.dev_record_file, config.batch_size)
lr = config.learning_rate
base_lr = 1
lr_warm_up_num = config.lr_warm_up_num
model = QANet(word_mat, char_mat).to(device)
ema = EMA(config.decay)
for name, param in model.named_parameters():
if param.requires_grad:
ema.register(name, param.data)
parameters = filter(lambda param: param.requires_grad, model.parameters())
optimizer = optim.Adam(lr=base_lr,
betas=(0.9, 0.999),
eps=1e-7,
weight_decay=5e-8,
params=parameters)
cr = lr / math.log2(lr_warm_up_num)
scheduler = optim.lr_scheduler.LambdaLR(
optimizer,
lr_lambda=lambda ee: cr * math.log2(ee + 1)
if ee < lr_warm_up_num else lr)
best_f1 = 0
best_em = 0
patience = 0
unused = False
for iter in range(config.num_epoch):
train(model, optimizer, scheduler, train_dataset, dev_dataset,
dev_eval_file, iter, ema)
ema.assign(model)
metrics = test(model, dev_dataset, dev_eval_file,
(iter + 1) * len(train_dataset))
dev_f1 = metrics["f1"]
dev_em = metrics["exact_match"]
if dev_f1 < best_f1 and dev_em < best_em:
patience += 1
if patience > config.early_stop:
break
else:
patience = 0
best_f1 = max(best_f1, dev_f1)
best_em = max(best_em, dev_em)
fn = os.path.join(config.save_dir, "model.pt")
torch.save(model, fn)
ema.resume(model)
