MultiTaskModel/model_train.py at main · busesese/MultiTaskModel · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
# -*- coding: utf-8 -*-
# @Time : 2021-04-19 17:10
# @Author : WenYi
# @Contact : 1244058349@qq.com
# @Description : model train function

importtorch
fromtqdmimporttqdm
fromsklearn.metricsimportroc_auc_score

deftrain_model(model, train_loader, val_loader, epoch, loss_function, optimizer, path, early_stop):
"""
 pytorch model train function
 :param model: pytorch model
 :param train_loader: dataloader, train data loader
 :param val_loader: dataloader, val data loader
 :param epoch: int, number of iters
 :param loss_function: loss function of train model
 :param optimizer: pytorch optimizer
 :param path: save path
 :param early_stop: int, early stop number
 :return: None
 """
# use GPU
device=torch.device("cuda"iftorch.cuda.is_available() else"cpu")
model.to(device)

# 多少步内验证集的loss没有变小就提前停止
patience, eval_loss=0, 0

# train
foriinrange(epoch):
y_train_income_true= []
y_train_income_predict= []
y_train_marry_true= []
y_train_marry_predict= []
total_loss, count=0, 0
foridx, (x, y1, y2) intqdm(enumerate(train_loader), total=len(train_loader)):
x, y1, y2=x.to(device), y1.to(device), y2.to(device)
predict=model(x)
y_train_income_true+=list(y1.squeeze().cpu().numpy())
y_train_marry_true+=list(y2.squeeze().cpu().numpy())
y_train_income_predict+=list(predict[0].squeeze().cpu().detach().numpy())
y_train_marry_predict+=list(predict[1].squeeze().cpu().detach().numpy())
loss_1=loss_function(predict[0], y1.unsqueeze(1).float())
loss_2=loss_function(predict[1], y2.unsqueeze(1).float())
loss=loss_1+loss_2
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_loss+=float(loss)
count+=1
torch.save(model, path.format(i+1))
income_auc=roc_auc_score(y_train_income_true, y_train_income_predict)
marry_auc=roc_auc_score(y_train_marry_true, y_train_marry_predict)
print("Epoch %d train loss is %.3f, income auc is %.3f and marry auc is %.3f"% (i+1, total_loss/count,
income_auc, marry_auc))

# 验证
total_eval_loss=0
model.eval()
count_eval=0
y_val_income_true= []
y_val_marry_true= []
y_val_income_predict= []
y_val_marry_predict= []
foridx, (x, y1, y2) intqdm(enumerate(val_loader), total=len(val_loader)):
x, y1, y2=x.to(device), y1.to(device), y2.to(device)
predict=model(x)
y_val_income_true+=list(y1.squeeze().cpu().numpy())
y_val_marry_true+=list(y2.squeeze().cpu().numpy())
y_val_income_predict+=list(predict[0].squeeze().cpu().detach().numpy())
y_val_marry_predict+=list(predict[1].squeeze().cpu().detach().numpy())
loss_1=loss_function(predict[0], y1.unsqueeze(1).float())
loss_2=loss_function(predict[1], y2.unsqueeze(1).float())
loss=loss_1+loss_2
total_eval_loss+=float(loss)
count_eval+=1
income_auc=roc_auc_score(y_val_income_true, y_val_income_predict)
marry_auc=roc_auc_score(y_val_marry_true, y_val_marry_predict)
print("Epoch %d val loss is %.3f, income auc is %.3f and marry auc is %.3f"% (i+1,
total_eval_loss/count_eval,
income_auc, marry_auc))

# earl stopping
ifi==0:
eval_loss=total_eval_loss/count_eval
else:
iftotal_eval_loss/count_eval<eval_loss:
eval_loss=total_eval_loss/count_eval
else:
ifpatience<early_stop:
patience+=1
else:
print("val loss is not decrease in %d epoch and break training"%patience)
break