-
Notifications
You must be signed in to change notification settings - Fork 4
/
train_dual.py
103 lines (87 loc) · 4.66 KB
/
train_dual.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
import os
import time
from options.train_options import TrainOptions
from options.test_options import TestOptions
from models.denseloss_model import *
from util.util import *
from datasets.base_dataset import CustomDatasetDataLoader
from models.base_model import create_model
def train():
print('Initialize Parameters...')
options = TrainOptions()
opt = options.parse()
options.save_to_file()
print('Load data...')
train_data_loader = CustomDatasetDataLoader(opt)
opt.phase = 'test'
test_data_loader = CustomDatasetDataLoader(opt)
opt.phase = 'train'
data_loader_size = len(train_data_loader)
print('Construct Model...')
model = create_model(opt)
print(opt)
#print(opt.model)
model.train()
total_steps = 0
print('Start Training...')
for epoch in range(opt.start_epoch, opt.start_epoch + opt.num_epoch):
epoch_start_time = time.time()
epoch_steps = 0
batch_start_time = time.time()
for batch_idx, batch_data in enumerate(train_data_loader):
epoch_steps += opt.batch_size
total_steps += opt.batch_size
iter_start_time = time.time()
model.optimize(batch_data)
if batch_idx % opt.print_freq == 0:
batch_end_time = time.time()
now_size = opt.batch_size * (batch_idx+1)
print('Train Epoch: {} [{}/{} ({:.2f}%)] Time:{:.6f} \t{}'.format(epoch,
now_size, data_loader_size, now_size / data_loader_size * 100.0,
batch_end_time - batch_start_time,
model.generate_message(model.result_record)))
batch_start_time = time.time()
if total_steps % opt.print_val_freq == 0 and (not opt.dataset_type in ['sketchy', 'imagenet'] or opt.model == 'cls_model') :
val_start_time = time.time()
now_size = opt.batch_size * (batch_idx+1)
for i, batch_test_data in enumerate(test_data_loader):
model.test(batch_test_data, opt.retrieval_now)
if not opt.retrieval_now:
model.test_features = model.combine_features(model.test_features)
model.retrieval_evaluation(model.test_features, model.test_result_record['total']['loss_value'].avg, model.test_features['labels'])
val_end_time = time.time()
print('Validation Epoch: {} [{}/{} ({:.2f}%)] Time:{:.6f} \t{}'.format(epoch,
now_size, data_loader_size, now_size / data_loader_size * 100.0,
val_end_time - val_start_time,
model.generate_message(model.test_result_record)))
model.reset_test_features()
model.reset_test_records()
if total_steps % opt.save_latest_freq == 0:
print('Save Model at latest epoch {} total steps {}.'.format(epoch, total_steps))
model.save_model('total_{}'.format(total_steps))
model.save_model('latest', True)
model.reset_records()
if not opt.dataset_type in ['sketchy', 'imagenet'] or opt.model == 'cls_model':
for i, batch_test_data in enumerate(test_data_loader):
model.test(batch_test_data, opt.retrieval_now)
if not opt.retrieval_now:
model.test_features = model.combine_features(model.test_features)
model.retrieval_evaluation(model.test_features, model.test_result_record['total']['loss_value'].avg, model.test_features['labels'])
if epoch % opt.save_epoch_freq == 0:
print('Save Model at epoch {}.'.format(epoch))
model.save_model('epoch_{}'.format(epoch))
model.save_model('latest', True)
epoch_end_time = time.time()
print('End of epoch {} / {}, Time:{:.6f}, \t {}'.format(epoch, opt.start_epoch + opt.num_epoch,
epoch_end_time - epoch_start_time,
model.generate_message(model.test_result_record)))
model.reset_features()
model.reset_test_features()
model.reset_records()
model.reset_test_records()
if epoch % 2 == 0:
opt.query_what = "sketch"
else:
opt.query_what = "image"
train_data_loader = CustomDatasetDataLoader(opt)
train()