def train(args):
data_folder = args.get("data_folder")
save_folder = args.get("save_folder")
image_size = args.get("image_size")
max_epoch = args.get("max_epoch")
save_epoch = args.get("save_epoch") or max_epoch // 10
# Scale lr and steps_per_epoch accordingly.
# Make sure the total number of gradient evaluations is consistent.
n_gpu = args.get("n_gpu") or 1
batch_size = args.get("batch_size") or BATCH
equi_batch_size = max(n_gpu, 1) * batch_size
lr = args.get("lr") or LR
lr *= equi_batch_size
steps_per_epoch = args.get("steps_per_epoch") or 1000
steps_per_epoch /= equi_batch_size
image_steps = args.get("image_steps") or steps_per_epoch // 10
scalar_steps = args.get("scalar_steps")
if scalar_steps > 0:
scalar_steps = max(scalar_steps // equi_batch_size, 1)
else:
scalar_steps = 0 # merge scalar summary every epoch
# lr starts decreasing at half of max epoch
start_dec_epoch = max_epoch // 2
# stops when lr is 0.01 of its initial value
end_epoch = max_epoch - int((max_epoch - start_dec_epoch) * 0.01)
# adjust noise input range according to the input act
zmin, zmax = (0, 1) if args.get("act_input") == "identity" else (-1, 1)
if save_folder == None:
logger.auto_set_dir()
else:
logger.set_logger_dir(save_folder)
df = get_data(data_folder,
image_size,
zmin=zmin,
zmax=zmax,
batch=batch_size)
df = PrintData(df)
data = QueueInput(df)
SynTexTrainer(data, Style2PO(args), n_gpu).train_with_defaults(
callbacks=[
PeriodicTrigger(ModelSaver(), every_k_epochs=save_epoch),
PeriodicTrigger(ModelSaver(),
every_k_epochs=end_epoch), # save model at last
ScheduledHyperParamSetter('learning_rate', [(start_dec_epoch, lr),
(max_epoch, 0)],
interp="linear"),
PeriodicTrigger(VisualizeTestSet(data_folder, image_size),
every_k_epochs=max(1, max_epoch // 100)),
#MergeAllSummaries(period=scalar_steps), # scalar only, slowdown in training, use TCMalloc
MergeAllSummaries(period=image_steps, key="image_summaries"),
MergeAllSummaries(key="acti_summaries"),
],
max_epoch=end_epoch,
steps_per_epoch=steps_per_epoch,
session_init=None)
steps_per_epoch /= equi_batch_size
image_steps = args.get("image_steps") or steps_per_epoch // 10
scalar_steps = args.get("scalar_steps")
if scalar_steps > 0:
scalar_steps = max(scalar_steps // equi_batch_size, 1)
else:
scalar_steps = 0 # merge scalar summary every epoch
# lr starts decreasing at half of max epoch
start_dec_epoch = max_epoch // 2
# stops when lr is 0.01 of its initial value
end_epoch = max_epoch - int((max_epoch - start_dec_epoch) * 0.01)
# adjust noise input range according to the input act
zmin, zmax = (0, 1) if args.get("act") == "identity" else (-1, 1)
if save_folder == None:
logger.auto_set_dir()
else:
logger.set_logger_dir(save_folder)
df = get_data(data_folder, image_size, zmin=zmin, zmax=zmax)
df = PrintData(df)
data = QueueInput(df)
SynTexTrainer(data, AdaptiveSynTex(args), n_gpu).train_with_defaults(
callbacks=[
PeriodicTrigger(ModelSaver(), every_k_epochs=save_epoch),
PeriodicTrigger(ModelSaver(),
every_k_epochs=end_epoch), # save model at last
ScheduledHyperParamSetter('learning_rate', [(start_dec_epoch, lr),
(max_epoch, 0)],
interp="linear"),
extra_callbacks=[
MovingAverageSummary(),
ProgressBar(["EMA/cost", "EMA/loss/pos-dist",
"EMA/loss/neg-dist"]),
MergeAllSummaries(),
RunUpdateOps()
],
nr_tower=nr_gpu,
steps_per_epoch=1800,
max_epoch=30,
session_config=tf.ConfigProto(allow_soft_placement=True))
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--gpu', help='comma separated list of GPU(s) to use.')
parser.add_argument('--lmdb', help='lmdb path of the arod dataset.')
parser.add_argument('--txt', help='txt with predefined triplets.')
parser.add_argument('--load',
help='load checkpoint file with pretrained models.')
args = parser.parse_args()
logger.auto_set_dir(name="resNetTriplet50_1024Features")
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
NR_GPU = len(args.gpu.split(','))
config = get_config(args.lmdb, args.txt, NR_GPU)
if args.load:
config.session_init = SaverRestore(args.load)
SyncMultiGPUTrainer(config).train()