main.py 3.5 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253
  1. import argparse
  2. import os
  3. import tensorflow as tf
  4. tf.set_random_seed(19)
  5. from model import cyclegan
  6. parser = argparse.ArgumentParser(description='')
  7. parser.add_argument('--dataset_dir', dest='dataset_dir', default='horse2zebra', help='path of the dataset')
  8. parser.add_argument('--epoch', dest='epoch', type=int, default=200, help='# of epoch')
  9. parser.add_argument('--epoch_step', dest='epoch_step', type=int, default=100, help='# of epoch to decay lr')
  10. parser.add_argument('--batch_size', dest='batch_size', type=int, default=1, help='# images in batch')
  11. parser.add_argument('--train_size', dest='train_size', type=int, default=1e8, help='# images used to train')
  12. parser.add_argument('--load_size', dest='load_size', type=int, default=286, help='scale images to this size')
  13. parser.add_argument('--fine_size', dest='fine_size', type=int, default=256, help='then crop to this size')
  14. parser.add_argument('--ngf', dest='ngf', type=int, default=64, help='# of gen filters in first conv layer')
  15. parser.add_argument('--ndf', dest='ndf', type=int, default=64, help='# of discri filters in first conv layer')
  16. parser.add_argument('--input_nc', dest='input_nc', type=int, default=3, help='# of input image channels')
  17. parser.add_argument('--output_nc', dest='output_nc', type=int, default=3, help='# of output image channels')
  18. parser.add_argument('--lr', dest='lr', type=float, default=0.0002, help='initial learning rate for adam')
  19. parser.add_argument('--beta1', dest='beta1', type=float, default=0.5, help='momentum term of adam')
  20. parser.add_argument('--which_direction', dest='which_direction', default='AtoB', help='AtoB or BtoA')
  21. parser.add_argument('--phase', dest='phase', default='train', help='train, test')
  22. parser.add_argument('--save_freq', dest='save_freq', type=int, default=1000, help='save a model every save_freq iterations')
  23. parser.add_argument('--print_freq', dest='print_freq', type=int, default=100, help='print the debug information every print_freq iterations')
  24. parser.add_argument('--continue_train', dest='continue_train', type=bool, default=False, help='if continue training, load the latest model: 1: true, 0: false')
  25. parser.add_argument('--checkpoint_dir', dest='checkpoint_dir', default='./checkpoint', help='models are saved here')
  26. parser.add_argument('--sample_dir', dest='sample_dir', default='./sample', help='sample are saved here')
  27. parser.add_argument('--test_dir', dest='test_dir', default='./test', help='test sample are saved here')
  28. parser.add_argument('--L1_lambda', dest='L1_lambda', type=float, default=10.0, help='weight on L1 term in objective')
  29. parser.add_argument('--use_resnet', dest='use_resnet', type=bool, default=True, help='generation network using reidule block')
  30. parser.add_argument('--use_lsgan', dest='use_lsgan', type=bool, default=True, help='gan loss defined in lsgan')
  31. parser.add_argument('--max_size', dest='max_size', type=int, default=50, help='max size of image pool, 0 means do not use image pool')
  32. args = parser.parse_args()
  33. def main(_):
  34. if not os.path.exists(args.checkpoint_dir):
  35. os.makedirs(args.checkpoint_dir)
  36. if not os.path.exists(args.sample_dir):
  37. os.makedirs(args.sample_dir)
  38. if not os.path.exists(args.test_dir):
  39. os.makedirs(args.test_dir)
  40. tfconfig = tf.ConfigProto(allow_soft_placement=True)
  41. tfconfig.gpu_options.allow_growth = True
  42. with tf.Session(config=tfconfig) as sess:
  43. model = cyclegan(sess, args)
  44. model.train(args) if args.phase == 'train' \
  45. else model.test(args)
  46. if __name__ == '__main__':
  47. tf.app.run()