parser.add_argument('--mode', choices=['AtoB', 'BtoA'], default='AtoB')
parser.add_argument('-b', '--batch', type=int, default=1)
args = parser.parse_args()
if args.gpu:
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
BATCH = args.batch
if args.sample:
assert args.load
sample(args.data, args.load)
else:
logger.auto_set_dir()
data = QueueInput(get_data())
nr_tower = max(get_num_gpu(), 1)
if nr_tower == 1:
trainer = GANTrainer(data, Model())
else:
trainer = MultiGPUGANTrainer(nr_tower, data, Model())
trainer.train_with_defaults(
callbacks=[
PeriodicTrigger(ModelSaver(), every_k_epochs=3),
ScheduledHyperParamSetter('learning_rate', [(200, 1e-4)])
],
steps_per_epoch=data.size(),
max_epoch=300,
session_init=SaverRestore(args.load) if args.load else None)
# sample
pass
else:
# Set up configuration
# Set the logger directory
logger.auto_set_dir()
# SyncMultiGPUTrainer(config).train()
nr_tower = max(get_nr_gpu(), 1)
if nr_tower == 1:
trainer = SeparateGANTrainer(data_set,
model,
g_period=4,
d_period=1)
else:
trainer = MultiGPUGANTrainer(nr_tower, data_set, model)
trainer.train_with_defaults(
callbacks=[
# PeriodicTrigger(ModelSaver(), every_k_epochs=20),
ClipCallback(),
ScheduledHyperParamSetter('learning_rate', [(0, 2e-4),
(100, 1e-4),
(200, 2e-5),
(300, 1e-5),
(400, 2e-6),
(500, 1e-6)],
interp='linear'),
PeriodicTrigger(VisualizeRunner(), every_k_epochs=5),
],
session_init=SaverRestore(args.load) if args.load else None,
steps_per_epoch=data_set.size(),
def get_config():
return TrainConfig(
model=Model(),
dataflow=DCGAN.get_data(G.data),
callbacks=[
ModelSaver(),
StatMonitorParamSetter('learning_rate', 'measure',
lambda x: x * 0.5, 0, 10)
],
steps_per_epoch=500,
max_epoch=400,
)
if __name__ == '__main__':
args = DCGAN.get_args()
if args.sample:
DCGAN.sample(args.load, 'gen/conv4.3/output')
else:
assert args.data
logger.auto_set_dir()
config = get_config()
if args.load:
config.session_init = SaverRestore(args.load)
nr_gpu = get_nr_gpu()
config.nr_tower = max(get_nr_gpu(), 1)
if config.nr_tower == 1:
GANTrainer(config).train()
else:
MultiGPUGANTrainer(config).train()
initializer=1e-4,
trainable=False)
opt = tf.train.AdamOptimizer(lr, beta1=0.5, beta2=0.9)
return opt
if __name__ == '__main__':
args = DCGAN.get_args(default_batch=32, default_z_dim=64)
if args.sample:
DCGAN.sample(Model(), args.load, 'gen/conv4.3/output')
else:
logger.auto_set_dir()
input = QueueInput(DCGAN.get_data())
model = Model()
nr_tower = max(get_nr_gpu(), 1)
if nr_tower == 1:
trainer = GANTrainer(input, model)
else:
trainer = MultiGPUGANTrainer(nr_tower, input, model)
trainer.train_with_defaults(
callbacks=[
ModelSaver(),
StatMonitorParamSetter('learning_rate', 'measure',
lambda x: x * 0.5, 0, 10)
],
session_init=SaverRestore(args.load) if args.load else None,
steps_per_epoch=500,
max_epoch=400)