def train(args, logdir):
# model
print("####model")
model = Net1()
# dataflow
print("####dataflow")
df = Net1DataFlow(hp.Train1.data_path, hp.Train1.batch_size)
# set logger for event and model saver
print("####logger")
logger.set_logger_dir(logdir)
print("####session_conf")
session_conf = tf.ConfigProto(gpu_options=tf.GPUOptions(
allow_growth=True, ),
allow_soft_placement=True)
print("####train_conf")
train_conf = TrainConfig(
model=model,
data=QueueInput(df(n_prefetch=1000, n_thread=5)),
callbacks=[
ModelSaver(checkpoint_dir=logdir),
# TODO EvalCallback()
],
max_epoch=hp.Train1.num_epochs,
steps_per_epoch=hp.Train1.steps_per_epoch,
session_config=session_conf)
print("####ckpt")
ckpt = '{}/{}'.format(
logdir, args.ckpt) if args.ckpt else tf.train.latest_checkpoint(logdir)
if ckpt:
train_conf.session_init = SaverRestore(ckpt)
if args.gpu:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
train_conf.nr_tower = len(args.gpu.split(','))
print("####trainer")
trainer = SyncMultiGPUTrainerReplicated(hp.Train1.num_gpu)
print("####launch_train_with_config")
launch_train_with_config(train_conf, trainer=trainer)
def train(args, logdir):
# model
model = Net()
# dataflow
df = NetDataFlow(hp.train.data_path, hp.train.batch_size)
# set logger for event and model saver
logger.set_logger_dir(logdir)
session_conf = tf.ConfigProto(
gpu_options=tf.GPUOptions(
allow_growth=True,
),)
session_conf.gpu_options.per_process_gpu_memory_fraction = 0.45 # 占用GPU90%的显存
train_conf = TrainConfig(
model=model,
data=QueueInput(df(n_prefetch=1000, n_thread=4)),
callbacks=[
ModelSaver(checkpoint_dir=logdir),
# TODO EvalCallback()
],
max_epoch=hp.train.num_epochs,
steps_per_epoch=hp.train.steps_per_epoch,
# session_config=session_conf
)
ckpt = '{}/{}'.format(logdir, args.ckpt) if args.ckpt else tf.train.latest_checkpoint(logdir)
if ckpt:
train_conf.session_init = SaverRestore(ckpt)
if args.gpu:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
train_conf.nr_tower = len(args.gpu.split(','))
trainer = SyncMultiGPUTrainerReplicated(hp.train.num_gpu)
launch_train_with_config(train_conf, trainer=trainer)
train_conf = TrainConfig(
model=model,
data=TFDatasetInput(dataset()),
callbacks=[
ModelSaver(checkpoint_dir=hp.logdir),
RunUpdateOps() # for batch norm, exponential moving average
# TODO GenerateCallback()
],
max_epoch=hp.train.num_epochs,
steps_per_epoch=hp.train.steps_per_epoch,
)
ckpt = '{}/{}'.format(
hp.logdir, ckpt) if ckpt else tf.train.latest_checkpoint(hp.logdir)
if ckpt:
train_conf.session_init = SaverRestore(ckpt)
if gpu is not None:
os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(map(str, gpu))
train_conf.nr_tower = len(gpu)
if hp.train.num_gpu <= 1:
trainer = SimpleTrainer()
else:
trainer = SyncMultiGPUTrainerReplicated(gpus=hp.train.num_gpu)
launch_train_with_config(train_conf, trainer=trainer)
if __name__ == '__main__':
fire.Fire(train)