104 lines
4.2 KiB
Python
104 lines
4.2 KiB
Python
|
# coding: UTF-8
|
|||
|
|
import numpy as np
|
|||
|
|
import torch
|
|||
|
|
import torch.nn as nn
|
|||
|
|
import torch.nn.functional as F
|
|||
|
|
from sklearn import metrics
|
|||
|
|
import time
|
|||
|
|
from utils import get_time_dif
|
|||
|
|
from transformers.optimization import AdamW
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
def train(config, model, train_iter, dev_iter, test_iter):
|
|||
|
|
start_time = time.time()
|
|||
|
|
model.train()
|
|||
|
|
param_optimizer = list(model.named_parameters())
|
|||
|
|
no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
|
|||
|
|
optimizer_grouped_parameters = [
|
|||
|
|
{'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01},
|
|||
|
|
{'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}]
|
|||
|
|
# optimizer = torch.optim.Adam(model.parameters(), lr=config.learning_rate)
|
|||
|
|
optimizer = AdamW(optimizer_grouped_parameters,
|
|||
|
|
lr=config.learning_rate,
|
|||
|
|
)
|
|||
|
|
total_batch = 0 # 记录进行到多少batch
|
|||
|
|
dev_best_loss = float('inf')
|
|||
|
|
last_improve = 0 # 记录上次验证集loss下降的batch数
|
|||
|
|
flag = False # 记录是否很久没有效果提升
|
|||
|
|
model.train()
|
|||
|
|
for epoch in range(config.num_epochs):
|
|||
|
|
print('Epoch [{}/{}]'.format(epoch + 1, config.num_epochs))
|
|||
|
|
for i, (trains, labels) in enumerate(train_iter):
|
|||
|
|
outputs = model(trains)
|
|||
|
|
model.zero_grad()
|
|||
|
|
loss = F.cross_entropy(outputs, labels)
|
|||
|
|
loss.backward()
|
|||
|
|
optimizer.step()
|
|||
|
|
if total_batch % 100 == 0:
|
|||
|
|
# 每多少轮输出在训练集和验证集上的效果
|
|||
|
|
true = labels.data.cpu()
|
|||
|
|
predic = torch.max(outputs.data, 1)[1].cpu()
|
|||
|
|
train_acc = metrics.accuracy_score(true, predic)
|
|||
|
|
dev_acc, dev_loss = evaluate(config, model, dev_iter)
|
|||
|
|
if dev_loss < dev_best_loss:
|
|||
|
|
dev_best_loss = dev_loss
|
|||
|
|
torch.save(model.state_dict(), config.save_path)
|
|||
|
|
improve = '*'
|
|||
|
|
last_improve = total_batch
|
|||
|
|
else:
|
|||
|
|
improve = ''
|
|||
|
|
time_dif = get_time_dif(start_time)
|
|||
|
|
msg = 'Iter: {0:>6}, Train Loss: {1:>5.2}, Train Acc: {2:>6.2%}, Val Loss: {3:>5.2}, Val Acc: {4:>6.2%}, Time: {5} {6}'
|
|||
|
|
print(msg.format(total_batch, loss.item(), train_acc, dev_loss, dev_acc, time_dif, improve))
|
|||
|
|
model.train()
|
|||
|
|
total_batch += 1
|
|||
|
|
if total_batch - last_improve > config.require_improvement:
|
|||
|
|
# 验证集loss超过1000batch没下降，结束训练
|
|||
|
|
print("No optimization for a long time, auto-stopping...")
|
|||
|
|
flag = True
|
|||
|
|
break
|
|||
|
|
if flag:
|
|||
|
|
break
|
|||
|
|
test(config, model, test_iter)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test(config, model, test_iter):
|
|||
|
|
# test
|
|||
|
|
model.load_state_dict(torch.load(config.save_path))
|
|||
|
|
model.eval()
|
|||
|
|
start_time = time.time()
|
|||
|
|
test_acc, test_loss, test_report, test_confusion = evaluate(config, model, test_iter, test=True)
|
|||
|
|
msg = 'Test Loss: {0:>5.2}, Test Acc: {1:>6.2%}'
|
|||
|
|
print(msg.format(test_loss, test_acc))
|
|||
|
|
print("Precision, Recall and F1-Score...")
|
|||
|
|
print(test_report)
|
|||
|
|
print("Confusion Matrix...")
|
|||
|
|
print(test_confusion)
|
|||
|
|
time_dif = get_time_dif(start_time)
|
|||
|
|
print("Time usage:", time_dif)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def evaluate(config, model, data_iter, test=False):
|
|||
|
|
model.eval()
|
|||
|
|
loss_total = 0
|
|||
|
|
predict_all = np.array([], dtype=int)
|
|||
|
|
labels_all = np.array([], dtype=int)
|
|||
|
|
with torch.no_grad():
|
|||
|
|
for texts, labels in data_iter:
|
|||
|
|
outputs = model(texts)
|
|||
|
|
loss = F.cross_entropy(outputs, labels)
|
|||
|
|
loss_total += loss
|
|||
|
|
labels = labels.data.cpu().numpy()
|
|||
|
|
predic = torch.max(outputs.data, 1)[1].cpu().numpy()
|
|||
|
|
labels_all = np.append(labels_all, labels)
|
|||
|
|
predict_all = np.append(predict_all, predic)
|
|||
|
|
|
|||
|
|
acc = metrics.accuracy_score(labels_all, predict_all)
|
|||
|
|
if test:
|
|||
|
|
report = metrics.classification_report(labels_all, predict_all, target_names=config.class_list, digits=4)
|
|||
|
|
confusion = metrics.confusion_matrix(labels_all, predict_all)
|
|||
|
|
return acc, loss_total / len(data_iter), report, confusion
|
|||
|
|
return acc, loss_total / len(data_iter)
|