/
train.py
executable file
·131 lines (95 loc) · 4.83 KB
/
train.py
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
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
import os
import torch
from torch.autograd import Variable
from utils.parameters_loader import ParametersLoader
from utils.save import Save
from dataset_init.dataset_init import DatasetInit
from models.model_init import ModelIinit
import configparser
import time
class Train(ParametersLoader):
def __init__(self, *args, **kwargs):
super(Train, self).__init__(*args, **kwargs)
self.config = config
self.dataset_init = DatasetInit(self.config)
self.model_init = ModelIinit(self.config)
self.save = Save()
def Train(self):
self.dataloaders = self.dataset_init.dataloaders
'''training process'''
n_correct = 0
accuracy_old = {"train": 0, "val": 0}
for epoch in range(self.model_params["num_epochs"]):
print("epoch:{}, time:{}".format(epoch, time.strftime("%X")))
type_id_accuracy = {"train": {}, "val": {}}
loss_average = {"train": 0, "val": 0}
accuracy = {"train": 0, "val": 0}
for phase in ["train", "val"]:
type_id_accuracy[phase] = {}
loss_average[phase] = 0
accuracy[phase] = 0
if phase == 'train':
'''Set model to training mode'''
self.model_init.model.train()
else:
'''Set model to evaluate mode'''
self.model_init.model.eval()
for next_data in self.dataloaders[phase]:
'''calculate gradients only over train phase'''
with torch.set_grad_enabled(phase == 'train'):
images = next_data[0]
labels = next_data[1]
length = next_data[2]
label_text = next_data[3]
cards_id = next_data[4]
types_id = next_data[5]
images = images.cuda()
length = length.cuda()
labels = labels.cuda()
preds = self.model_init.model(images)
preds_size = Variable(torch.IntTensor([preds.size(0)] *
self.model_params["batch_size"]))
length = torch.squeeze(length)
preds_size = preds_size.cuda()
probs, preds_copy = preds.max(2)
preds = preds.log_softmax(2)
cost = self.model_init.criterion(preds, labels, preds_size,
length) / self.model_params["batch_size"] # length - real target length
loss_average[phase] = loss_average[phase] + cost
if phase == "train":
self.model_init.model.zero_grad()
cost.backward()
self.model_init.optimizer.step()
'''evaluate accuracy'''
preds_copy = preds_copy.to("cpu")
preds_copy = preds_copy.transpose(1, 0).contiguous().view(-1)
sim_preds = self.model_init.converter.decode(preds_copy.data, preds_size.data, raw=False)
for sample, (logit, label, card_id, type_id) in enumerate(
zip(sim_preds, label_text, cards_id, types_id)):
if logit == label:
n_correct += 1
loss_average[phase] = loss_average[phase]/\
(len(self.dataloaders[phase])*self.model_params["batch_size"])
accuracy[phase] = n_correct/( len(self.dataloaders[phase])*self.model_params["batch_size"])
if phase == "val":
if accuracy["val"] > accuracy_old["val"]:
print("saving_model")
torch.save(
self.model_init.model.state_dict(),
os.path.join(self.path_to_model, 'crnn_number.pth'))
accuracy_old["val"] = accuracy["val"]
'''save metrics of loss and accuracy'''
metrics = {}
for param_name, param in zip(
["loss_average", "accuracy"],
[loss_average, accuracy]):
for phase_ in ["train", "val"]:
metrics[param_name + "_" + phase_] = [round(float(param[phase_]), 6)]
self.save.save_metrix(self.path_to_metrics, metrics)
n_correct = 0
CONFIG_FILENAME = 'number_recognition.ini'
config = configparser.ConfigParser()
config.read(CONFIG_FILENAME)
train = Train(config=config)
'''start model training'''
train.Train()