-
Notifications
You must be signed in to change notification settings - Fork 3
/
train.py
141 lines (115 loc) · 5.57 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
"""
# ==================================
# AUTHOR : Yan Li, Qiong Wang
# CREATE DATE : 03.13.2020
# Contact : liyanxian19@gmail.com
# ==================================
# Change History: None
# ==================================
"""
from __future__ import print_function
########## Set seed ##########
from numpy.random import seed
from tensorflow import set_random_seed
seed(1)
set_random_seed(2)
########## import python libs ##########
import argparse
import math
import os
from distutils.version import LooseVersion
########## Import third-party libs ##########
from tensorflow.python import keras
assert LooseVersion(keras.__version__) > LooseVersion('2.0.8')
from tensorflow.python.keras import backend as K
########## Import our libs ##########
from dataset import *
from model import llfnet
from utils import *
parser = argparse.ArgumentParser()
parser.add_argument('--batch_size', type=int, default=1)
parser.add_argument('--gpus', type=int, default=1)
parser.add_argument('--workers', type=int, default=8)
parser.add_argument('--data_root', type=str, default='./Data') # YOUR_DATA_ROOT
parser.add_argument('--dataset', type=str, default='Hand-designed_test') # YOUR_DATA_SET.txt in YOUR_DATA_ROOT
parser.add_argument('--data_size', type=str, default='Image') # Changed to 'Patch' if there exists CPU/GPU memory issues
parser.add_argument('--iterations', type=int, default=10)
parser.add_argument("--learning_rate", type=float, default=1e-4)
parser.add_argument('--model_path', type=str, default='./Model') # YOUR_MODEL_PATH
parser.add_argument('--move_path', type=str, default='LT') # LT: Left-right, Top-bottom
parser.add_argument('--model_infovis', type=bool, default=True)
parser.add_argument("--finetune", action='store_true')
config = parser.parse_args()
def train(train_generator, model_train, iter_n, max_iter,
val_generator=None, val_steps=None, logger=None):
init_lr = config.learning_rate
# Update learning rate
schedule = {"step": {"80": 5e-5, "160": 1e-5, "200": 1e-6}}
update_lr = update_learning_rate(schedule, init_lr, iter_n, max_iter)
if update_lr != init_lr:
init_lr = update_lr
K.set_value(model_train.optimizer.lr, init_lr)
logger.info('=> updated learning rate: {}'.format(K.get_value(model_train.optimizer.lr)))
steps_per_epoch = config.iterations
if val_generator is None:
model_train.fit_generator(train_generator, steps_per_epoch=steps_per_epoch,
epochs=iter_n+1, initial_epoch=iter_n, verbose=2, workers=config.workers)
else:
model_train.fit_generator(train_generator, steps_per_epoch=steps_per_epoch,
epochs=iter_n+1, initial_epoch=iter_n,
validation_data=val_generator, validation_steps=val_steps,
verbose=2, workers=config.workers)
model_weights = model_train.get_weights()
return model_weights
def run(hp, logger=None):
########## prepare dataset ##########
train_generator = prepare_train_generator(hp, logger=logger)
input_chns = int(math.ceil(len(hp["network"]["train_crop_seqs"]) / hp["network"]["sample"]))
# padding
if "train_pad" in hp["network"].keys():
pad_n_h = hp["network"]["train_pad"][0]
pad_n_w = hp["network"]["train_pad"][1]
input_shape = (hp["network"]["train_size"][0]+2*pad_n_h,
hp["network"]["train_size"][1]+2*pad_n_w, input_chns)
else:
input_shape = (hp["network"]["train_size"][0],
hp["network"]["train_size"][1], input_chns)
########## prepare model ##########
input_layer_names = hp["network"]["input_layer_candidates"]
llfnet_model = llfnet(input_layer_names, input_shape, config=config, hp=hp, phase="train", logger=logger)
# finetune the model?
if config.finetune:
logger.info("=> finetune from {}".format(config.model_path))
llfnet_model, iter_st = load_weight(llfnet_model, config.model_path)
else:
iter_st = 0
max_iter = hp["max_n_iter"]
########## start to train (patch-wise training) ##########
for iter_n in range(iter_st, max_iter):
logger.info("------------------ iteration {}----------------------".format(iter_n))
model_weights = train(train_generator,
llfnet_model,
iter_n,
max_iter)
save_path_file_new = '%s_iter%04d' % ('wlf_train', iter_n)
if not os.path.exists(os.path.join(config.model_path, save_path_file_new + '.hdf5')):
llfnet_model.save(os.path.join(config.model_path, save_path_file_new + '.hdf5'))
if __name__ == '__main__':
########## call logger ##########
logger = get_logger()
########## setup GPU ##########
if config.gpus > 1:
pass
else:
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' # disable printing tensorflow gpu info etc.
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
# load hyper-parameters
baseline = 1
density = 9
dataset_repr = "ca" # ma: micro-lens array (narrow), ca: camera array (wide)
json_file_name = "hyper-parameters_llf-net_b{}d{}_{}.json".format(baseline,
density,
dataset_repr)
hp = get_hp("manual", json_file_name)
run(hp=hp, logger=logger)