def train():
save_dir = "/home/FuDawei/NLP/SQUAD/unet/data/"
train_examples, dev_examples, opt = prepare_train(save_dir)
epoch = 30
batch_size = 32
model = Unet(opt=opt).to(device)
parameters = filter(lambda p: p.requires_grad, model.parameters())
optimizer = torch.optim.Adamax(parameters, lr=opt["lr"])
best_score, exact_scores, f1_scores = 0, [], []
count = 0
total_loss = 0
for ep in range(epoch):
model.train()
for batch_data in get_batch_data(train_examples, batch_size):
data = model.get_data(batch_data)
loss = model(data)
model.zero_grad()
optimizer.zero_grad()
loss.backward()
nn.utils.clip_grad_norm_(parameters, 10)
optimizer.step()
model.reset_parameters()
count += 1
# print(loss.item())
# Evaluate(dev_examples, model)
total_loss += loss.item()
if count % 100 == 0:
print(total_loss / 100)
total_loss = 0
# model.eval()
# Evaluate(dev_examples, model, opt)
if not opt["fix_word_embedding"]:
model.reset_parameters()
print(ep)
model.eval()
exact, f1 = Evaluate(dev_examples, model, opt)
exact_scores.append(exact)
f1_scores.append(f1)
if f1 > best_score:
best_score = f1
torch.save(model.state_dict(), save_dir + "best_model")
with open(save_dir + '_f1_scores.pkl', 'wb') as f:
pkl.dump(f1_scores, f)
with open(save_dir + '_exact_scores.pkl', 'wb') as f:
pkl.dump(exact_scores, f)
def train(args, x_train, y_train, x_valid, y_valid):
writer = SummaryWriter()
best_dice = 0
model = Unet().to(args.device)
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
bce_loss = torch.nn.BCELoss()
train_dataloader = load_dataset(x_train, y_train, args.batch_size, True)
valid_dataloader = load_dataset(x_valid, y_valid, args.batch_size, False)
result = {}
result['train/BCE'] = []
result['train/Dice'] = []
result['valid/BCE'] = []
result['valid/Dice'] = []
for epoch in range(args.epochs):
print('train step: epoch {}'.format(str(epoch+1).zfill(4)))
train_bce = []
train_dice = []
for inp_im, lab_im in tqdm(train_dataloader):
inp_im = inp_im.to(args.device)
lab_im = lab_im.to(args.device)
pred = model(inp_im)
bce = bce_loss(pred, lab_im)
dice = calc_dice(pred, lab_im)
train_bce.append(bce.item())
train_dice.append(dice)
model.zero_grad()
bce.backward()
optimizer.step()
result['train/BCE'].append(statistics.mean(train_bce))
result['train/Dice'].append(statistics.mean(train_dice))
writer.add_scalar('train/BinaryCrossEntropy', result['train/BCE'][-1], epoch+1)
writer.add_scalar('train/DiceScore', result['train/Dice'][-1], epoch+1)
print('BCE: {}, Dice: {}'.format(result['train/BCE'][-1], result['train/Dice'][-1]))
if (epoch+1) % 10 == 0 or (epoch+1) == 1:
with torch.no_grad():
print('valid step: epoch {}'.format(str(epoch+1).zfill(4)))
model.eval()
valid_bce = []
valid_dice = []
for inp_im, lab_im in tqdm(valid_dataloader):
inp_im = inp_im.to(args.device)
lab_im = lab_im.to(args.device)
pred = model(inp_im)
bce = bce_loss(pred, lab_im)
dice = calc_dice(pred, lab_im)
valid_bce.append(bce.item())
valid_dice.append(dice)
result['valid/BCE'].append(statistics.mean(valid_bce))
result['valid/Dice'].append(statistics.mean(valid_dice))
writer.add_scalar('valid/BinaryCrossEntropy', result['valid/BCE'][-1], epoch+1)
writer.add_scalar('valid/DiceScore', result['valid/Dice'][-1], epoch+1)
print('BCE: {}, Dice: {}'.format(result['valid/BCE'][-1], result['valid/Dice'][-1]))
if best_dice < result['valid/Dice'][-1]:
best_dice = result['valid/Dice'][-1]
best_model_name = os.path.join(args.save_model_path, f'best_model_{epoch + 1:04}.pth')
print('save model ==>> {}'.format(best_model_name))
torch.save(model.state_dict(), best_model_name)
