-
Notifications
You must be signed in to change notification settings - Fork 1
/
train.py
147 lines (122 loc) · 6.45 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
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
import argparse
import os
import time
from datetime import datetime
import numpy as np
import torch
import torch.nn.functional
import torch.utils.data
import torch.optim as optim
from torch.autograd import Variable
from dataset import Dataset
from model import Model
from evaluator import Evaluator
parser = argparse.ArgumentParser()
parser.add_argument('-d', '--data_dir', default='./data', help='directory to read LMDB files')
parser.add_argument('-l', '--logdir', default='./logs', help='directory to write logs')
parser.add_argument('-r', '--restore_checkpoint', default=None,
help='path to restore checkpoint, e.g. ./logs/model-100.tar')
parser.add_argument('-b', '--batch_size', default=32, type=int, help='Default 32')
parser.add_argument('-lr', '--learning_rate', default=1e-2, type=float, help='Default 1e-2')
parser.add_argument('-p', '--patience', default=100, type=int, help='Default 100, set -1 to train infinitely')
parser.add_argument('-ds', '--decay_steps', default=10000, type=int, help='Default 10000')
parser.add_argument('-dr', '--decay_rate', default=0.9, type=float, help='Default 0.9')
def _loss(length_logits, digits_logits, length_labels, digits_labels):
length_cross_entropy = torch.nn.functional.cross_entropy(length_logits, length_labels)
digit1_cross_entropy = torch.nn.functional.cross_entropy(digits_logits[0], digits_labels[0])
digit2_cross_entropy = torch.nn.functional.cross_entropy(digits_logits[1], digits_labels[1])
digit3_cross_entropy = torch.nn.functional.cross_entropy(digits_logits[2], digits_labels[2])
digit4_cross_entropy = torch.nn.functional.cross_entropy(digits_logits[3], digits_labels[3])
digit5_cross_entropy = torch.nn.functional.cross_entropy(digits_logits[4], digits_labels[4])
loss = length_cross_entropy + digit1_cross_entropy + digit2_cross_entropy + digit3_cross_entropy + digit4_cross_entropy + digit5_cross_entropy
return loss
def _adjust_learning_rate(optimizer, step, initial_lr, decay_steps, decay_rate):
lr = initial_lr * (decay_rate ** (step // decay_steps))
for param_group in optimizer.param_groups:
param_group['lr'] = lr
return lr
def _train(path_to_train_lmdb_dir, path_to_val_lmdb_dir, path_to_log_dir,
path_to_restore_checkpoint_file, training_options):
batch_size = training_options['batch_size']
initial_learning_rate = training_options['learning_rate']
initial_patience = training_options['patience']
num_steps_to_show_loss = 100
num_steps_to_check = 1000
step = 0
patience = initial_patience
best_accuracy = 0.0
duration = 0.0
model = Model()
model.cuda()
if path_to_restore_checkpoint_file is not None:
assert os.path.isfile(path_to_restore_checkpoint_file), '%s not found' % path_to_restore_checkpoint_file
step = model.load(path_to_restore_checkpoint_file)
print 'Model restored from file: %s' % path_to_restore_checkpoint_file
train_loader = torch.utils.data.DataLoader(Dataset(path_to_train_lmdb_dir),
batch_size=batch_size, shuffle=True,
num_workers=2, pin_memory=True)
evaluator = Evaluator(path_to_val_lmdb_dir)
optimizer = optim.SGD(model.parameters(), lr=initial_learning_rate, momentum=0.9, weight_decay=0.0005)
path_to_losses_npy_file = os.path.join(path_to_log_dir, 'losses.npy')
if os.path.isfile(path_to_losses_npy_file):
losses = np.load(path_to_losses_npy_file)
else:
losses = np.empty([0], dtype=np.float32)
while True:
for batch_idx, (images, length_labels, digits_labels) in enumerate(train_loader):
start_time = time.time()
images, length_labels, digits_labels = (Variable(images.cuda()),
Variable(length_labels.cuda()),
[Variable(digit_labels.cuda()) for digit_labels in digits_labels])
length_logits, digits_logits = model.train()(images)
loss = _loss(length_logits, digits_logits, length_labels, digits_labels)
learning_rate = _adjust_learning_rate(optimizer, step=step, initial_lr=initial_learning_rate,
decay_steps=training_options['decay_steps'],
decay_rate=training_options['decay_rate'])
optimizer.zero_grad()
loss.backward()
optimizer.step()
step += 1
duration += time.time() - start_time
if step % num_steps_to_show_loss == 0:
examples_per_sec = batch_size * num_steps_to_show_loss / duration
duration = 0.0
print '=> %s: step %d, loss = %f, learning_rate = %f (%.1f examples/sec)' % (
datetime.now(), step, loss.data[0], learning_rate, examples_per_sec)
if step % num_steps_to_check != 0:
continue
losses = np.append(losses, loss.cpu().data.numpy())
np.save(path_to_losses_npy_file, losses)
print '=> Evaluating on validation dataset...'
accuracy = evaluator.evaluate(model)
print '==> accuracy = %f, best accuracy %f' % (accuracy, best_accuracy)
if accuracy > best_accuracy:
path_to_checkpoint_file = model.save(path_to_log_dir, step=step)
print '=> Model saved to file: %s' % path_to_checkpoint_file
patience = initial_patience
best_accuracy = accuracy
else:
patience -= 1
print '=> patience = %d' % patience
if patience == 0:
return
def main(args):
path_to_train_lmdb_dir = os.path.join(args.data_dir, 'train.lmdb')
path_to_val_lmdb_dir = os.path.join(args.data_dir, 'val.lmdb')
path_to_log_dir = args.logdir
path_to_restore_checkpoint_file = args.restore_checkpoint
training_options = {
'batch_size': args.batch_size,
'learning_rate': args.learning_rate,
'patience': args.patience,
'decay_steps': args.decay_steps,
'decay_rate': args.decay_rate
}
if not os.path.exists(path_to_log_dir):
os.makedirs(path_to_log_dir)
print 'Start training'
_train(path_to_train_lmdb_dir, path_to_val_lmdb_dir, path_to_log_dir,
path_to_restore_checkpoint_file, training_options)
print 'Done'
if __name__ == '__main__':
main(parser.parse_args())