diff --git a/build_data.py b/build_data.py index fb4f2eb..217ae93 100644 --- a/build_data.py +++ b/build_data.py @@ -3,114 +3,116 @@ import os try: - from os import scandir + from os import scandir except ImportError: - # Python 2 polyfill module - from scandir import scandir - + # Python 2 polyfill module + from scandir import scandir -FLAGS = tf.flags.FLAGS +FLAGS = tf.compat.v1.flags.FLAGS -tf.flags.DEFINE_string('X_input_dir', 'data/apple2orange/trainA', +tf.compat.v1.flags.DEFINE_string('X_input_dir', 'data/apple2orange/trainA', 'X input directory, default: data/apple2orange/trainA') -tf.flags.DEFINE_string('Y_input_dir', 'data/apple2orange/trainB', +tf.compat.v1.flags.DEFINE_string('Y_input_dir', 'data/apple2orange/trainB', 'Y input directory, default: data/apple2orange/trainB') -tf.flags.DEFINE_string('X_output_file', 'data/tfrecords/apple.tfrecords', +tf.compat.v1.flags.DEFINE_string('X_output_file', 'data/tfrecords/apple.tfrecords', 'X output tfrecords file, default: data/tfrecords/apple.tfrecords') -tf.flags.DEFINE_string('Y_output_file', 'data/tfrecords/orange.tfrecords', +tf.compat.v1.flags.DEFINE_string('Y_output_file', 'data/tfrecords/orange.tfrecords', 'Y output tfrecords file, default: data/tfrecords/orange.tfrecords') def data_reader(input_dir, shuffle=True): - """Read images from input_dir then shuffle them + """Read images from input_dir then shuffle them Args: input_dir: string, path of input dir, e.g., /path/to/dir Returns: file_paths: list of strings """ - file_paths = [] + file_paths = [] - for img_file in scandir(input_dir): - if img_file.name.endswith('.jpg') and img_file.is_file(): - file_paths.append(img_file.path) + for img_file in scandir(input_dir): + if img_file.name.endswith('.jpg') and img_file.is_file(): + file_paths.append(img_file.path) - if shuffle: - # Shuffle the ordering of all image files in order to guarantee - # random ordering of the images with respect to label in the - # saved TFRecord files. Make the randomization repeatable. - shuffled_index = list(range(len(file_paths))) - random.seed(12345) - random.shuffle(shuffled_index) + if shuffle: + # Shuffle the ordering of all image files in order to guarantee + # random ordering of the images with respect to label in the + # saved TFRecord files. Make the randomization repeatable. + shuffled_index = list(range(len(file_paths))) + random.seed(12345) + random.shuffle(shuffled_index) - file_paths = [file_paths[i] for i in shuffled_index] + file_paths = [file_paths[i] for i in shuffled_index] - return file_paths + return file_paths def _int64_feature(value): - """Wrapper for inserting int64 features into Example proto.""" - if not isinstance(value, list): - value = [value] - return tf.train.Feature(int64_list=tf.train.Int64List(value=value)) + """Wrapper for inserting int64 features into Example proto.""" + if not isinstance(value, list): + value = [value] + return tf.train.Feature(int64_list=tf.train.Int64List(value=value)) def _bytes_feature(value): - """Wrapper for inserting bytes features into Example proto.""" - return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value])) + """Wrapper for inserting bytes features into Example proto.""" + return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value])) def _convert_to_example(file_path, image_buffer): - """Build an Example proto for an example. + """Build an Example proto for an example. Args: file_path: string, path to an image file, e.g., '/path/to/example.JPG' image_buffer: string, JPEG encoding of RGB image Returns: Example proto """ - file_name = file_path.split('/')[-1] + file_name = file_path.split('/')[-1] - example = tf.train.Example(features=tf.train.Features(feature={ - 'image/file_name': _bytes_feature(tf.compat.as_bytes(os.path.basename(file_name))), - 'image/encoded_image': _bytes_feature((image_buffer)) + example = tf.train.Example(features=tf.train.Features(feature={ + 'image/file_name': _bytes_feature(tf.compat.as_bytes(os.path.basename(file_name))), + 'image/encoded_image': _bytes_feature((image_buffer)) })) - return example + return example + def data_writer(input_dir, output_file): - """Write data to tfrecords + """Write data to tfrecords """ - file_paths = data_reader(input_dir) + file_paths = data_reader(input_dir) + + # create tfrecords dir if not exists + output_dir = os.path.dirname(output_file) + try: + os.makedirs(output_dir) + except os.error as e: + pass - # create tfrecords dir if not exists - output_dir = os.path.dirname(output_file) - try: - os.makedirs(output_dir) - except os.error as e: - pass + images_num = len(file_paths) - images_num = len(file_paths) + # dump to tfrecords file + writer = tf.compat.v1.python_io.TFRecordWriter(output_file) - # dump to tfrecords file - writer = tf.python_io.TFRecordWriter(output_file) + for i in range(len(file_paths)): + file_path = file_paths[i] - for i in range(len(file_paths)): - file_path = file_paths[i] + with tf.compat.v1.gfile.FastGFile(file_path, 'rb') as f: + image_data = f.read() - with tf.gfile.FastGFile(file_path, 'rb') as f: - image_data = f.read() + example = _convert_to_example(file_path, image_data) + writer.write(example.SerializeToString()) - example = _convert_to_example(file_path, image_data) - writer.write(example.SerializeToString()) + if i % 500 == 0: + print("Processed {}/{}.".format(i, images_num)) + print("Done.") + writer.close() - if i % 500 == 0: - print("Processed {}/{}.".format(i, images_num)) - print("Done.") - writer.close() def main(unused_argv): - print("Convert X data to tfrecords...") - data_writer(FLAGS.X_input_dir, FLAGS.X_output_file) - print("Convert Y data to tfrecords...") - data_writer(FLAGS.Y_input_dir, FLAGS.Y_output_file) + print("Convert X data to tfrecords...") + data_writer(FLAGS.X_input_dir, FLAGS.X_output_file) + print("Convert Y data to tfrecords...") + data_writer(FLAGS.Y_input_dir, FLAGS.Y_output_file) + if __name__ == '__main__': - tf.app.run() + tf.compat.v1.app.run() diff --git a/discriminator.py b/discriminator.py index 162ec6d..9f7b261 100644 --- a/discriminator.py +++ b/discriminator.py @@ -1,39 +1,40 @@ import tensorflow as tf import ops + class Discriminator: - def __init__(self, name, is_training, norm='instance', use_sigmoid=False): - self.name = name - self.is_training = is_training - self.norm = norm - self.reuse = False - self.use_sigmoid = use_sigmoid + def __init__(self, name, is_training, norm='instance', use_sigmoid=False): + self.name = name + self.is_training = is_training + self.norm = norm + self.reuse = False + self.use_sigmoid = use_sigmoid - def __call__(self, input): - """ + def __call__(self, input): + """ Args: input: batch_size x image_size x image_size x 3 Returns: output: 4D tensor batch_size x out_size x out_size x 1 (default 1x5x5x1) filled with 0.9 if real, 0.0 if fake """ - with tf.variable_scope(self.name): - # convolution layers - C64 = ops.Ck(input, 64, reuse=self.reuse, norm=None, - is_training=self.is_training, name='C64') # (?, w/2, h/2, 64) - C128 = ops.Ck(C64, 128, reuse=self.reuse, norm=self.norm, - is_training=self.is_training, name='C128') # (?, w/4, h/4, 128) - C256 = ops.Ck(C128, 256, reuse=self.reuse, norm=self.norm, - is_training=self.is_training, name='C256') # (?, w/8, h/8, 256) - C512 = ops.Ck(C256, 512,reuse=self.reuse, norm=self.norm, - is_training=self.is_training, name='C512') # (?, w/16, h/16, 512) + with tf.compat.v1.variable_scope(self.name): + # convolution layers + C64 = ops.Ck(input, 64, reuse=self.reuse, norm=None, + is_training=self.is_training, name='C64') # (?, w/2, h/2, 64) + C128 = ops.Ck(C64, 128, reuse=self.reuse, norm=self.norm, + is_training=self.is_training, name='C128') # (?, w/4, h/4, 128) + C256 = ops.Ck(C128, 256, reuse=self.reuse, norm=self.norm, + is_training=self.is_training, name='C256') # (?, w/8, h/8, 256) + C512 = ops.Ck(C256, 512, reuse=self.reuse, norm=self.norm, + is_training=self.is_training, name='C512') # (?, w/16, h/16, 512) - # apply a convolution to produce a 1 dimensional output (1 channel?) - # use_sigmoid = False if use_lsgan = True - output = ops.last_conv(C512, reuse=self.reuse, - use_sigmoid=self.use_sigmoid, name='output') # (?, w/16, h/16, 1) + # apply a convolution to produce a 1 dimensional output (1 channel?) + # use_sigmoid = False if use_lsgan = True + output = ops.last_conv(C512, reuse=self.reuse, + use_sigmoid=self.use_sigmoid, name='output') # (?, w/16, h/16, 1) - self.reuse = True - self.variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope=self.name) + self.reuse = True + self.variables = tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.TRAINABLE_VARIABLES, scope=self.name) - return output + return output diff --git a/export_graph.py b/export_graph.py index 7f94a6f..c7bc2da 100644 --- a/export_graph.py +++ b/export_graph.py @@ -13,48 +13,51 @@ from model import CycleGAN import utils -FLAGS = tf.flags.FLAGS +FLAGS = tf.compat.v1.flags.FLAGS -tf.flags.DEFINE_string('checkpoint_dir', '', 'checkpoints directory path') -tf.flags.DEFINE_string('XtoY_model', 'apple2orange.pb', 'XtoY model name, default: apple2orange.pb') -tf.flags.DEFINE_string('YtoX_model', 'orange2apple.pb', 'YtoX model name, default: orange2apple.pb') -tf.flags.DEFINE_integer('image_size', '256', 'image size, default: 256') -tf.flags.DEFINE_integer('ngf', 64, +tf.compat.v1.flags.DEFINE_string('checkpoint_dir', '', 'checkpoints directory path') +tf.compat.v1.flags.DEFINE_string('XtoY_model', 'apple2orange.pb', 'XtoY model name, default: apple2orange.pb') +tf.compat.v1.flags.DEFINE_string('YtoX_model', 'orange2apple.pb', 'YtoX model name, default: orange2apple.pb') +tf.compat.v1.flags.DEFINE_integer('image_size', '256', 'image size, default: 256') +tf.compat.v1.flags.DEFINE_integer('ngf', 64, 'number of gen filters in first conv layer, default: 64') -tf.flags.DEFINE_string('norm', 'instance', +tf.compat.v1.flags.DEFINE_string('norm', 'instance', '[instance, batch] use instance norm or batch norm, default: instance') + def export_graph(model_name, XtoY=True): - graph = tf.Graph() + graph = tf.Graph() + + with graph.as_default(): + cycle_gan = CycleGAN(ngf=FLAGS.ngf, norm=FLAGS.norm, image_size=FLAGS.image_size) - with graph.as_default(): - cycle_gan = CycleGAN(ngf=FLAGS.ngf, norm=FLAGS.norm, image_size=FLAGS.image_size) + input_image = tf.compat.v1.placeholder(tf.float32, shape=[FLAGS.image_size, FLAGS.image_size, 3], name='input_image') + cycle_gan.model() + if XtoY: + output_image = cycle_gan.G.sample(tf.expand_dims(input_image, 0)) + else: + output_image = cycle_gan.F.sample(tf.expand_dims(input_image, 0)) - input_image = tf.placeholder(tf.float32, shape=[FLAGS.image_size, FLAGS.image_size, 3], name='input_image') - cycle_gan.model() - if XtoY: - output_image = cycle_gan.G.sample(tf.expand_dims(input_image, 0)) - else: - output_image = cycle_gan.F.sample(tf.expand_dims(input_image, 0)) + output_image = tf.identity(output_image, name='output_image') + restore_saver = tf.compat.v1.train.Saver() + export_saver = tf.compat.v1.train.Saver() - output_image = tf.identity(output_image, name='output_image') - restore_saver = tf.train.Saver() - export_saver = tf.train.Saver() + with tf.compat.v1.Session(graph=graph) as sess: + sess.run(tf.compat.v1.global_variables_initializer()) + latest_ckpt = tf.train.latest_checkpoint(FLAGS.checkpoint_dir) + restore_saver.restore(sess, latest_ckpt) + output_graph_def = tf.compat.v1.graph_util.convert_variables_to_constants( + sess, graph.as_graph_def(), [output_image.op.name]) - with tf.Session(graph=graph) as sess: - sess.run(tf.global_variables_initializer()) - latest_ckpt = tf.train.latest_checkpoint(FLAGS.checkpoint_dir) - restore_saver.restore(sess, latest_ckpt) - output_graph_def = tf.graph_util.convert_variables_to_constants( - sess, graph.as_graph_def(), [output_image.op.name]) + tf.compat.v1.train.write_graph(output_graph_def, 'pretrained', model_name, as_text=False) - tf.train.write_graph(output_graph_def, 'pretrained', model_name, as_text=False) def main(unused_argv): - print('Export XtoY model...') - export_graph(FLAGS.XtoY_model, XtoY=True) - print('Export YtoX model...') - export_graph(FLAGS.YtoX_model, XtoY=False) + print('Export XtoY model...') + export_graph(FLAGS.XtoY_model, XtoY=True) + print('Export YtoX model...') + export_graph(FLAGS.YtoX_model, XtoY=False) + if __name__ == '__main__': - tf.app.run() + tf.compat.v1.app.run() diff --git a/generator.py b/generator.py index ecb715b..5e2e510 100644 --- a/generator.py +++ b/generator.py @@ -2,56 +2,57 @@ import ops import utils + class Generator: - def __init__(self, name, is_training, ngf=64, norm='instance', image_size=128): - self.name = name - self.reuse = False - self.ngf = ngf - self.norm = norm - self.is_training = is_training - self.image_size = image_size - - def __call__(self, input): - """ + def __init__(self, name, is_training, ngf=64, norm='instance', image_size=128): + self.name = name + self.reuse = False + self.ngf = ngf + self.norm = norm + self.is_training = is_training + self.image_size = image_size + + def __call__(self, input): + """ Args: input: batch_size x width x height x 3 Returns: output: same size as input """ - with tf.variable_scope(self.name): - # conv layers - c7s1_32 = ops.c7s1_k(input, self.ngf, is_training=self.is_training, norm=self.norm, - reuse=self.reuse, name='c7s1_32') # (?, w, h, 32) - d64 = ops.dk(c7s1_32, 2*self.ngf, is_training=self.is_training, norm=self.norm, - reuse=self.reuse, name='d64') # (?, w/2, h/2, 64) - d128 = ops.dk(d64, 4*self.ngf, is_training=self.is_training, norm=self.norm, - reuse=self.reuse, name='d128') # (?, w/4, h/4, 128) - - if self.image_size <= 128: - # use 6 residual blocks for 128x128 images - res_output = ops.n_res_blocks(d128, reuse=self.reuse, n=6) # (?, w/4, h/4, 128) - else: - # 9 blocks for higher resolution - res_output = ops.n_res_blocks(d128, reuse=self.reuse, n=9) # (?, w/4, h/4, 128) - - # fractional-strided convolution - u64 = ops.uk(res_output, 2*self.ngf, is_training=self.is_training, norm=self.norm, - reuse=self.reuse, name='u64') # (?, w/2, h/2, 64) - u32 = ops.uk(u64, self.ngf, is_training=self.is_training, norm=self.norm, - reuse=self.reuse, name='u32', output_size=self.image_size) # (?, w, h, 32) - - # conv layer - # Note: the paper said that ReLU and _norm were used - # but actually tanh was used and no _norm here - output = ops.c7s1_k(u32, 3, norm=None, - activation='tanh', reuse=self.reuse, name='output') # (?, w, h, 3) - # set reuse=True for next call - self.reuse = True - self.variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope=self.name) - - return output - - def sample(self, input): - image = utils.batch_convert2int(self.__call__(input)) - image = tf.image.encode_jpeg(tf.squeeze(image, [0])) - return image + with tf.compat.v1.variable_scope(self.name): + # conv layers + c7s1_32 = ops.c7s1_k(input, self.ngf, is_training=self.is_training, norm=self.norm, + reuse=self.reuse, name='c7s1_32') # (?, w, h, 32) + d64 = ops.dk(c7s1_32, 2 * self.ngf, is_training=self.is_training, norm=self.norm, + reuse=self.reuse, name='d64') # (?, w/2, h/2, 64) + d128 = ops.dk(d64, 4 * self.ngf, is_training=self.is_training, norm=self.norm, + reuse=self.reuse, name='d128') # (?, w/4, h/4, 128) + + if self.image_size <= 128: + # use 6 residual blocks for 128x128 images + res_output = ops.n_res_blocks(d128, reuse=self.reuse, n=6) # (?, w/4, h/4, 128) + else: + # 9 blocks for higher resolution + res_output = ops.n_res_blocks(d128, reuse=self.reuse, n=9) # (?, w/4, h/4, 128) + + # fractional-strided convolution + u64 = ops.uk(res_output, 2 * self.ngf, is_training=self.is_training, norm=self.norm, + reuse=self.reuse, name='u64') # (?, w/2, h/2, 64) + u32 = ops.uk(u64, self.ngf, is_training=self.is_training, norm=self.norm, + reuse=self.reuse, name='u32', output_size=self.image_size) # (?, w, h, 32) + + # conv layer + # Note: the paper said that ReLU and _norm were used + # but actually tanh was used and no _norm here + output = ops.c7s1_k(u32, 3, norm=None, + activation='tanh', reuse=self.reuse, name='output') # (?, w, h, 3) + # set reuse=True for next call + self.reuse = True + self.variables = tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.TRAINABLE_VARIABLES, scope=self.name) + + return output + + def sample(self, input): + image = utils.batch_convert2int(self.__call__(input)) + image = tf.image.encode_jpeg(tf.squeeze(image, [0])) + return image diff --git a/inference.py b/inference.py index 490c763..fd1ee9f 100644 --- a/inference.py +++ b/inference.py @@ -11,39 +11,42 @@ from model import CycleGAN import utils -FLAGS = tf.flags.FLAGS +FLAGS = tf.compat.v1.flags.FLAGS + +tf.compat.v1.flags.DEFINE_string('model', '', 'model path (.pb)') +tf.compat.v1.flags.DEFINE_string('input', 'input_sample.jpg', 'input image path (.jpg)') +tf.compat.v1.flags.DEFINE_string('output', 'output_sample.jpg', 'output image path (.jpg)') +tf.compat.v1.flags.DEFINE_integer('image_size', '256', 'image size, default: 256') -tf.flags.DEFINE_string('model', '', 'model path (.pb)') -tf.flags.DEFINE_string('input', 'input_sample.jpg', 'input image path (.jpg)') -tf.flags.DEFINE_string('output', 'output_sample.jpg', 'output image path (.jpg)') -tf.flags.DEFINE_integer('image_size', '256', 'image size, default: 256') def inference(): - graph = tf.Graph() - - with graph.as_default(): - with tf.gfile.FastGFile(FLAGS.input, 'rb') as f: - image_data = f.read() - input_image = tf.image.decode_jpeg(image_data, channels=3) - input_image = tf.image.resize_images(input_image, size=(FLAGS.image_size, FLAGS.image_size)) - input_image = utils.convert2float(input_image) - input_image.set_shape([FLAGS.image_size, FLAGS.image_size, 3]) - - with tf.gfile.FastGFile(FLAGS.model, 'rb') as model_file: - graph_def = tf.GraphDef() - graph_def.ParseFromString(model_file.read()) - [output_image] = tf.import_graph_def(graph_def, - input_map={'input_image': input_image}, - return_elements=['output_image:0'], - name='output') - - with tf.Session(graph=graph) as sess: - generated = output_image.eval() - with open(FLAGS.output, 'wb') as f: - f.write(generated) + graph = tf.Graph() + + with graph.as_default(): + with tf.compat.v1.gfile.FastGFile(FLAGS.input, 'rb') as f: + image_data = f.read() + input_image = tf.image.decode_jpeg(image_data, channels=3) + input_image = tf.image.resize(input_image, size=(FLAGS.image_size, FLAGS.image_size)) + input_image = utils.convert2float(input_image) + input_image.set_shape([FLAGS.image_size, FLAGS.image_size, 3]) + + with tf.compat.v1.gfile.FastGFile(FLAGS.model, 'rb') as model_file: + graph_def = tf.compat.v1.GraphDef() + graph_def.ParseFromString(model_file.read()) + [output_image] = tf.import_graph_def(graph_def, + input_map={'input_image': input_image}, + return_elements=['output_image:0'], + name='output') + + with tf.compat.v1.Session(graph=graph) as sess: + generated = output_image.eval() + with open(FLAGS.output, 'wb') as f: + f.write(generated) + def main(unused_argv): - inference() + inference() + if __name__ == '__main__': - tf.app.run() + tf.compat.v1.app.run() diff --git a/model.py b/model.py index 98a9af1..cd75571 100644 --- a/model.py +++ b/model.py @@ -7,21 +7,22 @@ REAL_LABEL = 0.9 + class CycleGAN: - def __init__(self, - X_train_file='', - Y_train_file='', - batch_size=1, - image_size=256, - use_lsgan=True, - norm='instance', - lambda1=10, - lambda2=10, - learning_rate=2e-4, - beta1=0.5, - ngf=64 - ): - """ + def __init__(self, + X_train_file='', + Y_train_file='', + batch_size=1, + image_size=256, + use_lsgan=True, + norm='instance', + lambda1=10, + lambda2=10, + learning_rate=2e-4, + beta1=0.5, + ngf=64 + ): + """ Args: X_train_file: string, X tfrecords file for training Y_train_file: string Y tfrecords file for training @@ -35,112 +36,112 @@ def __init__(self, beta1: float, momentum term of Adam ngf: number of gen filters in first conv layer """ - self.lambda1 = lambda1 - self.lambda2 = lambda2 - self.use_lsgan = use_lsgan - use_sigmoid = not use_lsgan - self.batch_size = batch_size - self.image_size = image_size - self.learning_rate = learning_rate - self.beta1 = beta1 - self.X_train_file = X_train_file - self.Y_train_file = Y_train_file - - self.is_training = tf.placeholder_with_default(True, shape=[], name='is_training') - - self.G = Generator('G', self.is_training, ngf=ngf, norm=norm, image_size=image_size) - self.D_Y = Discriminator('D_Y', - self.is_training, norm=norm, use_sigmoid=use_sigmoid) - self.F = Generator('F', self.is_training, ngf=ngf, norm=norm, image_size=image_size) - self.D_X = Discriminator('D_X', - self.is_training, norm=norm, use_sigmoid=use_sigmoid) - - self.fake_x = tf.placeholder(tf.float32, - shape=[batch_size, image_size, image_size, 3]) - self.fake_y = tf.placeholder(tf.float32, - shape=[batch_size, image_size, image_size, 3]) - - def model(self): - X_reader = Reader(self.X_train_file, name='X', - image_size=self.image_size, batch_size=self.batch_size) - Y_reader = Reader(self.Y_train_file, name='Y', - image_size=self.image_size, batch_size=self.batch_size) - - x = X_reader.feed() - y = Y_reader.feed() - - cycle_loss = self.cycle_consistency_loss(self.G, self.F, x, y) - - # X -> Y - fake_y = self.G(x) - G_gan_loss = self.generator_loss(self.D_Y, fake_y, use_lsgan=self.use_lsgan) - G_loss = G_gan_loss + cycle_loss - D_Y_loss = self.discriminator_loss(self.D_Y, y, self.fake_y, use_lsgan=self.use_lsgan) - - # Y -> X - fake_x = self.F(y) - F_gan_loss = self.generator_loss(self.D_X, fake_x, use_lsgan=self.use_lsgan) - F_loss = F_gan_loss + cycle_loss - D_X_loss = self.discriminator_loss(self.D_X, x, self.fake_x, use_lsgan=self.use_lsgan) - - # summary - tf.summary.histogram('D_Y/true', self.D_Y(y)) - tf.summary.histogram('D_Y/fake', self.D_Y(self.G(x))) - tf.summary.histogram('D_X/true', self.D_X(x)) - tf.summary.histogram('D_X/fake', self.D_X(self.F(y))) - - tf.summary.scalar('loss/G', G_gan_loss) - tf.summary.scalar('loss/D_Y', D_Y_loss) - tf.summary.scalar('loss/F', F_gan_loss) - tf.summary.scalar('loss/D_X', D_X_loss) - tf.summary.scalar('loss/cycle', cycle_loss) - - tf.summary.image('X/generated', utils.batch_convert2int(self.G(x))) - tf.summary.image('X/reconstruction', utils.batch_convert2int(self.F(self.G(x)))) - tf.summary.image('Y/generated', utils.batch_convert2int(self.F(y))) - tf.summary.image('Y/reconstruction', utils.batch_convert2int(self.G(self.F(y)))) - - return G_loss, D_Y_loss, F_loss, D_X_loss, fake_y, fake_x - - def optimize(self, G_loss, D_Y_loss, F_loss, D_X_loss): - def make_optimizer(loss, variables, name='Adam'): - """ Adam optimizer with learning rate 0.0002 for the first 100k steps (~100 epochs) + self.lambda1 = lambda1 + self.lambda2 = lambda2 + self.use_lsgan = use_lsgan + use_sigmoid = not use_lsgan + self.batch_size = batch_size + self.image_size = image_size + self.learning_rate = learning_rate + self.beta1 = beta1 + self.X_train_file = X_train_file + self.Y_train_file = Y_train_file + + self.is_training = tf.compat.v1.placeholder_with_default(True, shape=[], name='is_training') + + self.G = Generator('G', self.is_training, ngf=ngf, norm=norm, image_size=image_size) + self.D_Y = Discriminator('D_Y', + self.is_training, norm=norm, use_sigmoid=use_sigmoid) + self.F = Generator('F', self.is_training, ngf=ngf, norm=norm, image_size=image_size) + self.D_X = Discriminator('D_X', + self.is_training, norm=norm, use_sigmoid=use_sigmoid) + + self.fake_x = tf.compat.v1.placeholder(tf.float32, + shape=[batch_size, image_size, image_size, 3]) + self.fake_y = tf.compat.v1.placeholder(tf.float32, + shape=[batch_size, image_size, image_size, 3]) + + def model(self): + X_reader = Reader(self.X_train_file, name='X', + image_size=self.image_size, batch_size=self.batch_size) + Y_reader = Reader(self.Y_train_file, name='Y', + image_size=self.image_size, batch_size=self.batch_size) + + x = X_reader.feed() + y = Y_reader.feed() + + cycle_loss = self.cycle_consistency_loss(self.G, self.F, x, y) + + # X -> Y + fake_y = self.G(x) + G_gan_loss = self.generator_loss(self.D_Y, fake_y, use_lsgan=self.use_lsgan) + G_loss = G_gan_loss + cycle_loss + D_Y_loss = self.discriminator_loss(self.D_Y, y, self.fake_y, use_lsgan=self.use_lsgan) + + # Y -> X + fake_x = self.F(y) + F_gan_loss = self.generator_loss(self.D_X, fake_x, use_lsgan=self.use_lsgan) + F_loss = F_gan_loss + cycle_loss + D_X_loss = self.discriminator_loss(self.D_X, x, self.fake_x, use_lsgan=self.use_lsgan) + + # summary + tf.summary.histogram('D_Y/true', self.D_Y(y)) + tf.summary.histogram('D_Y/fake', self.D_Y(self.G(x))) + tf.summary.histogram('D_X/true', self.D_X(x)) + tf.summary.histogram('D_X/fake', self.D_X(self.F(y))) + + tf.summary.scalar('loss/G', G_gan_loss) + tf.summary.scalar('loss/D_Y', D_Y_loss) + tf.summary.scalar('loss/F', F_gan_loss) + tf.summary.scalar('loss/D_X', D_X_loss) + tf.summary.scalar('loss/cycle', cycle_loss) + + tf.summary.image('X/generated', utils.batch_convert2int(self.G(x))) + tf.summary.image('X/reconstruction', utils.batch_convert2int(self.F(self.G(x)))) + tf.summary.image('Y/generated', utils.batch_convert2int(self.F(y))) + tf.summary.image('Y/reconstruction', utils.batch_convert2int(self.G(self.F(y)))) + + return G_loss, D_Y_loss, F_loss, D_X_loss, fake_y, fake_x + + def optimize(self, G_loss, D_Y_loss, F_loss, D_X_loss): + def make_optimizer(loss, variables, name='Adam'): + """ Adam optimizer with learning rate 0.0002 for the first 100k steps (~100 epochs) and a linearly decaying rate that goes to zero over the next 100k steps """ - global_step = tf.Variable(0, trainable=False) - starter_learning_rate = self.learning_rate - end_learning_rate = 0.0 - start_decay_step = 100000 - decay_steps = 100000 - beta1 = self.beta1 - learning_rate = ( - tf.where( - tf.greater_equal(global_step, start_decay_step), - tf.train.polynomial_decay(starter_learning_rate, global_step-start_decay_step, - decay_steps, end_learning_rate, - power=1.0), - starter_learning_rate - ) - - ) - tf.summary.scalar('learning_rate/{}'.format(name), learning_rate) - - learning_step = ( - tf.train.AdamOptimizer(learning_rate, beta1=beta1, name=name) - .minimize(loss, global_step=global_step, var_list=variables) - ) - return learning_step - - G_optimizer = make_optimizer(G_loss, self.G.variables, name='Adam_G') - D_Y_optimizer = make_optimizer(D_Y_loss, self.D_Y.variables, name='Adam_D_Y') - F_optimizer = make_optimizer(F_loss, self.F.variables, name='Adam_F') - D_X_optimizer = make_optimizer(D_X_loss, self.D_X.variables, name='Adam_D_X') - - with tf.control_dependencies([G_optimizer, D_Y_optimizer, F_optimizer, D_X_optimizer]): - return tf.no_op(name='optimizers') - - def discriminator_loss(self, D, y, fake_y, use_lsgan=True): - """ Note: default: D(y).shape == (batch_size,5,5,1), + global_step = tf.Variable(0, trainable=False) + starter_learning_rate = self.learning_rate + end_learning_rate = 0.0 + start_decay_step = 100000 + decay_steps = 100000 + beta1 = self.beta1 + learning_rate = ( + tf.where( + tf.greater_equal(global_step, start_decay_step), + tf.compat.v1.train.polynomial_decay(starter_learning_rate, global_step - start_decay_step, + decay_steps, end_learning_rate, + power=1.0), + starter_learning_rate + ) + + ) + tf.summary.scalar('learning_rate/{}'.format(name), learning_rate) + + learning_step = ( + tf.compat.v1.train.AdamOptimizer(learning_rate, beta1=beta1, name=name) + .minimize(loss, global_step=global_step, var_list=variables) + ) + return learning_step + + G_optimizer = make_optimizer(G_loss, self.G.variables, name='Adam_G') + D_Y_optimizer = make_optimizer(D_Y_loss, self.D_Y.variables, name='Adam_D_Y') + F_optimizer = make_optimizer(F_loss, self.F.variables, name='Adam_F') + D_X_optimizer = make_optimizer(D_X_loss, self.D_X.variables, name='Adam_D_X') + + with tf.control_dependencies([G_optimizer, D_Y_optimizer, F_optimizer, D_X_optimizer]): + return tf.no_op(name='optimizers') + + def discriminator_loss(self, D, y, fake_y, use_lsgan=True): + """ Note: default: D(y).shape == (batch_size,5,5,1), fake_buffer_size=50, batch_size=1 Args: G: generator object @@ -149,32 +150,32 @@ def discriminator_loss(self, D, y, fake_y, use_lsgan=True): Returns: loss: scalar """ - if use_lsgan: - # use mean squared error - error_real = tf.reduce_mean(tf.squared_difference(D(y), REAL_LABEL)) - error_fake = tf.reduce_mean(tf.square(D(fake_y))) - else: - # use cross entropy - error_real = -tf.reduce_mean(ops.safe_log(D(y))) - error_fake = -tf.reduce_mean(ops.safe_log(1-D(fake_y))) - loss = (error_real + error_fake) / 2 - return loss - - def generator_loss(self, D, fake_y, use_lsgan=True): - """ fool discriminator into believing that G(x) is real + if use_lsgan: + # use mean squared error + error_real = tf.reduce_mean(tf.math.squared_difference(D(y), REAL_LABEL)) + error_fake = tf.reduce_mean(tf.square(D(fake_y))) + else: + # use cross entropy + error_real = -tf.reduce_mean(ops.safe_log(D(y))) + error_fake = -tf.reduce_mean(ops.safe_log(1 - D(fake_y))) + loss = (error_real + error_fake) / 2 + return loss + + def generator_loss(self, D, fake_y, use_lsgan=True): + """ fool discriminator into believing that G(x) is real """ - if use_lsgan: - # use mean squared error - loss = tf.reduce_mean(tf.squared_difference(D(fake_y), REAL_LABEL)) - else: - # heuristic, non-saturating loss - loss = -tf.reduce_mean(ops.safe_log(D(fake_y))) / 2 - return loss - - def cycle_consistency_loss(self, G, F, x, y): - """ cycle consistency loss (L1 norm) + if use_lsgan: + # use mean squared error + loss = tf.reduce_mean(tf.math.squared_difference(D(fake_y), REAL_LABEL)) + else: + # heuristic, non-saturating loss + loss = -tf.reduce_mean(ops.safe_log(D(fake_y))) / 2 + return loss + + def cycle_consistency_loss(self, G, F, x, y): + """ cycle consistency loss (L1 norm) """ - forward_loss = tf.reduce_mean(tf.abs(F(G(x))-x)) - backward_loss = tf.reduce_mean(tf.abs(G(F(y))-y)) - loss = self.lambda1*forward_loss + self.lambda2*backward_loss - return loss + forward_loss = tf.reduce_mean(tf.abs(F(G(x)) - x)) + backward_loss = tf.reduce_mean(tf.abs(G(F(y)) - y)) + loss = self.lambda1 * forward_loss + self.lambda2 * backward_loss + return loss diff --git a/ops.py b/ops.py index 9becd24..79f87f0 100644 --- a/ops.py +++ b/ops.py @@ -1,9 +1,10 @@ import tensorflow as tf + ## Layers: follow the naming convention used in the original paper ### Generator layers def c7s1_k(input, k, reuse=False, norm='instance', activation='relu', is_training=True, name='c7s1_k'): - """ A 7x7 Convolution-BatchNorm-ReLU layer with k filters and stride 1 + """ A 7x7 Convolution-BatchNorm-ReLU layer with k filters and stride 1 Args: input: 4D tensor k: integer, number of filters (output depth) @@ -16,24 +17,25 @@ def c7s1_k(input, k, reuse=False, norm='instance', activation='relu', is_trainin Returns: 4D tensor """ - with tf.variable_scope(name, reuse=reuse): - weights = _weights("weights", - shape=[7, 7, input.get_shape()[3], k]) + with tf.compat.v1.variable_scope(name, reuse=reuse): + weights = _weights("weights", + shape=[7, 7, input.get_shape()[3], k]) - padded = tf.pad(input, [[0,0],[3,3],[3,3],[0,0]], 'REFLECT') - conv = tf.nn.conv2d(padded, weights, - strides=[1, 1, 1, 1], padding='VALID') + padded = tf.pad(input, [[0, 0], [3, 3], [3, 3], [0, 0]], 'REFLECT') + conv = tf.nn.conv2d(padded, weights, + strides=[1, 1, 1, 1], padding='VALID') - normalized = _norm(conv, is_training, norm) + normalized = _norm(conv, is_training, norm) + + if activation == 'relu': + output = tf.nn.relu(normalized) + if activation == 'tanh': + output = tf.nn.tanh(normalized) + return output - if activation == 'relu': - output = tf.nn.relu(normalized) - if activation == 'tanh': - output = tf.nn.tanh(normalized) - return output def dk(input, k, reuse=False, norm='instance', is_training=True, name=None): - """ A 3x3 Convolution-BatchNorm-ReLU layer with k filters and stride 2 + """ A 3x3 Convolution-BatchNorm-ReLU layer with k filters and stride 2 Args: input: 4D tensor k: integer, number of filters (output depth) @@ -45,18 +47,19 @@ def dk(input, k, reuse=False, norm='instance', is_training=True, name=None): Returns: 4D tensor """ - with tf.variable_scope(name, reuse=reuse): - weights = _weights("weights", - shape=[3, 3, input.get_shape()[3], k]) - - conv = tf.nn.conv2d(input, weights, - strides=[1, 2, 2, 1], padding='SAME') - normalized = _norm(conv, is_training, norm) - output = tf.nn.relu(normalized) - return output + with tf.compat.v1.variable_scope(name, reuse=reuse): + weights = _weights("weights", + shape=[3, 3, input.get_shape()[3], k]) -def Rk(input, k, reuse=False, norm='instance', is_training=True, name=None): - """ A residual block that contains two 3x3 convolutional layers + conv = tf.nn.conv2d(input, weights, + strides=[1, 2, 2, 1], padding='SAME') + normalized = _norm(conv, is_training, norm) + output = tf.nn.relu(normalized) + return output + + +def Rk(input, k, reuse=False, norm='instance', is_training=True, name=None): + """ A residual block that contains two 3x3 convolutional layers with the same number of filters on both layer Args: input: 4D Tensor @@ -66,36 +69,38 @@ def Rk(input, k, reuse=False, norm='instance', is_training=True, name=None): Returns: 4D tensor (same shape as input) """ - with tf.variable_scope(name, reuse=reuse): - with tf.variable_scope('layer1', reuse=reuse): - weights1 = _weights("weights1", - shape=[3, 3, input.get_shape()[3], k]) - padded1 = tf.pad(input, [[0,0],[1,1],[1,1],[0,0]], 'REFLECT') - conv1 = tf.nn.conv2d(padded1, weights1, - strides=[1, 1, 1, 1], padding='VALID') - normalized1 = _norm(conv1, is_training, norm) - relu1 = tf.nn.relu(normalized1) - - with tf.variable_scope('layer2', reuse=reuse): - weights2 = _weights("weights2", - shape=[3, 3, relu1.get_shape()[3], k]) - - padded2 = tf.pad(relu1, [[0,0],[1,1],[1,1],[0,0]], 'REFLECT') - conv2 = tf.nn.conv2d(padded2, weights2, - strides=[1, 1, 1, 1], padding='VALID') - normalized2 = _norm(conv2, is_training, norm) - output = input+normalized2 - return output + with tf.compat.v1.variable_scope(name, reuse=reuse): + with tf.compat.v1.variable_scope('layer1', reuse=reuse): + weights1 = _weights("weights1", + shape=[3, 3, input.get_shape()[3], k]) + padded1 = tf.pad(input, [[0, 0], [1, 1], [1, 1], [0, 0]], 'REFLECT') + conv1 = tf.nn.conv2d(padded1, weights1, + strides=[1, 1, 1, 1], padding='VALID') + normalized1 = _norm(conv1, is_training, norm) + relu1 = tf.nn.relu(normalized1) + + with tf.compat.v1.variable_scope('layer2', reuse=reuse): + weights2 = _weights("weights2", + shape=[3, 3, relu1.get_shape()[3], k]) + + padded2 = tf.pad(relu1, [[0, 0], [1, 1], [1, 1], [0, 0]], 'REFLECT') + conv2 = tf.nn.conv2d(padded2, weights2, + strides=[1, 1, 1, 1], padding='VALID') + normalized2 = _norm(conv2, is_training, norm) + output = input + normalized2 + return output + def n_res_blocks(input, reuse, norm='instance', is_training=True, n=6): - depth = input.get_shape()[3] - for i in range(1,n+1): - output = Rk(input, depth, reuse, norm, is_training, 'R{}_{}'.format(depth, i)) - input = output - return output + depth = input.get_shape()[3] + for i in range(1, n + 1): + output = Rk(input, depth, reuse, norm, is_training, 'R{}_{}'.format(depth, i)) + input = output + return output + def uk(input, k, reuse=False, norm='instance', is_training=True, name=None, output_size=None): - """ A 3x3 fractional-strided-Convolution-BatchNorm-ReLU layer + """ A 3x3 fractional-strided-Convolution-BatchNorm-ReLU layer with k filters, stride 1/2 Args: input: 4D tensor @@ -108,25 +113,26 @@ def uk(input, k, reuse=False, norm='instance', is_training=True, name=None, outp Returns: 4D tensor """ - with tf.variable_scope(name, reuse=reuse): - input_shape = input.get_shape().as_list() - - weights = _weights("weights", - shape=[3, 3, k, input_shape[3]]) - - if not output_size: - output_size = input_shape[1]*2 - output_shape = [input_shape[0], output_size, output_size, k] - fsconv = tf.nn.conv2d_transpose(input, weights, - output_shape=output_shape, - strides=[1, 2, 2, 1], padding='SAME') - normalized = _norm(fsconv, is_training, norm) - output = tf.nn.relu(normalized) - return output + with tf.compat.v1.variable_scope(name, reuse=reuse): + input_shape = input.get_shape().as_list() + + weights = _weights("weights", + shape=[3, 3, k, input_shape[3]]) + + if not output_size: + output_size = input_shape[1] * 2 + output_shape = [input_shape[0], output_size, output_size, k] + fsconv = tf.nn.conv2d_transpose(input, weights, + output_shape=output_shape, + strides=[1, 2, 2, 1], padding='SAME') + normalized = _norm(fsconv, is_training, norm) + output = tf.nn.relu(normalized) + return output + ### Discriminator layers def Ck(input, k, slope=0.2, stride=2, reuse=False, norm='instance', is_training=True, name=None): - """ A 4x4 Convolution-BatchNorm-LeakyReLU layer with k filters and stride 2 + """ A 4x4 Convolution-BatchNorm-LeakyReLU layer with k filters and stride 2 Args: input: 4D tensor k: integer, number of filters (output depth) @@ -139,19 +145,20 @@ def Ck(input, k, slope=0.2, stride=2, reuse=False, norm='instance', is_training= Returns: 4D tensor """ - with tf.variable_scope(name, reuse=reuse): - weights = _weights("weights", - shape=[4, 4, input.get_shape()[3], k]) + with tf.compat.v1.variable_scope(name, reuse=reuse): + weights = _weights("weights", + shape=[4, 4, input.get_shape()[3], k]) - conv = tf.nn.conv2d(input, weights, - strides=[1, stride, stride, 1], padding='SAME') + conv = tf.nn.conv2d(input, weights, + strides=[1, stride, stride, 1], padding='SAME') + + normalized = _norm(conv, is_training, norm) + output = _leaky_relu(normalized, slope) + return output - normalized = _norm(conv, is_training, norm) - output = _leaky_relu(normalized, slope) - return output def last_conv(input, reuse=False, use_sigmoid=False, name=None): - """ Last convolutional layer of discriminator network + """ Last convolutional layer of discriminator network (1 filter with size 4x4, stride 1) Args: input: 4D tensor @@ -159,21 +166,22 @@ def last_conv(input, reuse=False, use_sigmoid=False, name=None): use_sigmoid: boolean (False if use lsgan) name: string, e.g. 'C64' """ - with tf.variable_scope(name, reuse=reuse): - weights = _weights("weights", - shape=[4, 4, input.get_shape()[3], 1]) - biases = _biases("biases", [1]) - - conv = tf.nn.conv2d(input, weights, - strides=[1, 1, 1, 1], padding='SAME') - output = conv + biases - if use_sigmoid: - output = tf.sigmoid(output) - return output + with tf.compat.v1.variable_scope(name, reuse=reuse): + weights = _weights("weights", + shape=[4, 4, input.get_shape()[3], 1]) + biases = _biases("biases", [1]) + + conv = tf.nn.conv2d(input, weights, + strides=[1, 1, 1, 1], padding='SAME') + output = conv + biases + if use_sigmoid: + output = tf.sigmoid(output) + return output + ### Helpers def _weights(name, shape, mean=0.0, stddev=0.02): - """ Helper to create an initialized Variable + """ Helper to create an initialized Variable Args: name: name of the variable shape: list of ints @@ -182,53 +190,59 @@ def _weights(name, shape, mean=0.0, stddev=0.02): Returns: A trainable variable """ - var = tf.get_variable( - name, shape, - initializer=tf.random_normal_initializer( - mean=mean, stddev=stddev, dtype=tf.float32)) - return var + var = tf.compat.v1.get_variable( + name, shape, + initializer=tf.compat.v1.random_normal_initializer( + mean=mean, stddev=stddev, dtype=tf.float32)) + return var + def _biases(name, shape, constant=0.0): - """ Helper to create an initialized Bias with constant + """ Helper to create an initialized Bias with constant """ - return tf.get_variable(name, shape, - initializer=tf.constant_initializer(constant)) + return tf.compat.v1.get_variable(name, shape, + initializer=tf.constant_initializer(constant)) + def _leaky_relu(input, slope): - return tf.maximum(slope*input, input) + return tf.maximum(slope * input, input) + def _norm(input, is_training, norm='instance'): - """ Use Instance Normalization or Batch Normalization or None + """ Use Instance Normalization or Batch Normalization or None """ - if norm == 'instance': - return _instance_norm(input) - elif norm == 'batch': - return _batch_norm(input, is_training) - else: - return input + if norm == 'instance': + return _instance_norm(input) + elif norm == 'batch': + return _batch_norm(input, is_training) + else: + return input + def _batch_norm(input, is_training): - """ Batch Normalization + """ Batch Normalization """ - with tf.variable_scope("batch_norm"): - return tf.contrib.layers.batch_norm(input, - decay=0.9, - scale=True, - updates_collections=None, - is_training=is_training) + with tf.compat.v1.variable_scope("batch_norm"): + return tf.compat.v1.layers.batch_normalization(input, + decay=0.9, + scale=True, + updates_collections=None, + is_training=is_training) + def _instance_norm(input): - """ Instance Normalization + """ Instance Normalization """ - with tf.variable_scope("instance_norm"): - depth = input.get_shape()[3] - scale = _weights("scale", [depth], mean=1.0) - offset = _biases("offset", [depth]) - mean, variance = tf.nn.moments(input, axes=[1,2], keep_dims=True) - epsilon = 1e-5 - inv = tf.rsqrt(variance + epsilon) - normalized = (input-mean)*inv - return scale*normalized + offset + with tf.compat.v1.variable_scope("instance_norm"): + depth = input.get_shape()[3] + scale = _weights("scale", [depth], mean=1.0) + offset = _biases("offset", [depth]) + mean, variance = tf.nn.moments(input, axes=[1, 2], keepdims=True) + epsilon = 1e-5 + inv = tf.math.rsqrt(variance + epsilon) + normalized = (input - mean) * inv + return scale * normalized + offset + def safe_log(x, eps=1e-12): - return tf.log(x + eps) + return tf.math.log(x + eps) diff --git a/reader.py b/reader.py index f61e285..ce822e2 100644 --- a/reader.py +++ b/reader.py @@ -1,93 +1,96 @@ import tensorflow as tf import utils + class Reader(): - def __init__(self, tfrecords_file, image_size=256, - min_queue_examples=1000, batch_size=1, num_threads=8, name=''): - """ + def __init__(self, tfrecords_file, image_size=256, + min_queue_examples=1000, batch_size=1, num_threads=8, name=''): + """ Args: tfrecords_file: string, tfrecords file path min_queue_examples: integer, minimum number of samples to retain in the queue that provides of batches of examples batch_size: integer, number of images per batch num_threads: integer, number of preprocess threads """ - self.tfrecords_file = tfrecords_file - self.image_size = image_size - self.min_queue_examples = min_queue_examples - self.batch_size = batch_size - self.num_threads = num_threads - self.reader = tf.TFRecordReader() - self.name = name + self.tfrecords_file = tfrecords_file + self.image_size = image_size + self.min_queue_examples = min_queue_examples + self.batch_size = batch_size + self.num_threads = num_threads + self.reader = tf.compat.v1.TFRecordReader() + self.name = name - def feed(self): - """ + def feed(self): + """ Returns: images: 4D tensor [batch_size, image_width, image_height, image_depth] """ - with tf.name_scope(self.name): - filename_queue = tf.train.string_input_producer([self.tfrecords_file]) - reader = tf.TFRecordReader() + with tf.name_scope(self.name): + filename_queue = tf.compat.v1.train.string_input_producer([self.tfrecords_file]) + reader = tf.compat.v1.TFRecordReader() - _, serialized_example = self.reader.read(filename_queue) - features = tf.parse_single_example( - serialized_example, - features={ - 'image/file_name': tf.FixedLenFeature([], tf.string), - 'image/encoded_image': tf.FixedLenFeature([], tf.string), - }) + _, serialized_example = self.reader.read(filename_queue) + features = tf.compat.v1.parse_single_example( + serialized_example, + features={ + 'image/file_name': tf.compat.v1.FixedLenFeature([], tf.string), + 'image/encoded_image': tf.compat.v1.FixedLenFeature([], tf.string), + }) - image_buffer = features['image/encoded_image'] - image = tf.image.decode_jpeg(image_buffer, channels=3) - image = self._preprocess(image) - images = tf.train.shuffle_batch( - [image], batch_size=self.batch_size, num_threads=self.num_threads, - capacity=self.min_queue_examples + 3*self.batch_size, - min_after_dequeue=self.min_queue_examples - ) + image_buffer = features['image/encoded_image'] + image = tf.image.decode_jpeg(image_buffer, channels=3) + image = self._preprocess(image) + images = tf.compat.v1.train.shuffle_batch( + [image], batch_size=self.batch_size, num_threads=self.num_threads, + capacity=self.min_queue_examples + 3 * self.batch_size, + min_after_dequeue=self.min_queue_examples + ) - tf.summary.image('_input', images) - return images + tf.summary.image('_input', images) + return images + + def _preprocess(self, image): + image = tf.image.resize(image, size=(self.image_size, self.image_size)) + image = utils.convert2float(image) + image.set_shape([self.image_size, self.image_size, 3]) + return image - def _preprocess(self, image): - image = tf.image.resize_images(image, size=(self.image_size, self.image_size)) - image = utils.convert2float(image) - image.set_shape([self.image_size, self.image_size, 3]) - return image def test_reader(): - TRAIN_FILE_1 = 'data/tfrecords/apple.tfrecords' - TRAIN_FILE_2 = 'data/tfrecords/orange.tfrecords' + TRAIN_FILE_1 = 'data/tfrecords/apple.tfrecords' + TRAIN_FILE_2 = 'data/tfrecords/orange.tfrecords' + + with tf.Graph().as_default(): + reader1 = Reader(TRAIN_FILE_1, batch_size=2) + reader2 = Reader(TRAIN_FILE_2, batch_size=2) + images_op1 = reader1.feed() + images_op2 = reader2.feed() - with tf.Graph().as_default(): - reader1 = Reader(TRAIN_FILE_1, batch_size=2) - reader2 = Reader(TRAIN_FILE_2, batch_size=2) - images_op1 = reader1.feed() - images_op2 = reader2.feed() + sess = tf.Session() + init = tf.global_variables_initializer() + sess.run(init) - sess = tf.Session() - init = tf.global_variables_initializer() - sess.run(init) + coord = tf.train.Coordinator() + threads = tf.train.start_queue_runners(sess=sess, coord=coord) - coord = tf.train.Coordinator() - threads = tf.train.start_queue_runners(sess=sess, coord=coord) + try: + step = 0 + while not coord.should_stop(): + batch_images1, batch_images2 = sess.run([images_op1, images_op2]) + print("image shape: {}".format(batch_images1)) + print("image shape: {}".format(batch_images2)) + print("=" * 10) + step += 1 + except KeyboardInterrupt: + print('Interrupted') + coord.request_stop() + except Exception as e: + coord.request_stop(e) + finally: + # When done, ask the threads to stop. + coord.request_stop() + coord.join(threads) - try: - step = 0 - while not coord.should_stop(): - batch_images1, batch_images2 = sess.run([images_op1, images_op2]) - print("image shape: {}".format(batch_images1)) - print("image shape: {}".format(batch_images2)) - print("="*10) - step += 1 - except KeyboardInterrupt: - print('Interrupted') - coord.request_stop() - except Exception as e: - coord.request_stop(e) - finally: - # When done, ask the threads to stop. - coord.request_stop() - coord.join(threads) if __name__ == '__main__': - test_reader() + test_reader() diff --git a/train.py b/train.py index 0b0cc27..a042564 100644 --- a/train.py +++ b/train.py @@ -1,135 +1,136 @@ import tensorflow as tf from model import CycleGAN -from reader import Reader from datetime import datetime import os import logging from utils import ImagePool -FLAGS = tf.flags.FLAGS +FLAGS = tf.compat.v1.flags.FLAGS -tf.flags.DEFINE_integer('batch_size', 1, 'batch size, default: 1') -tf.flags.DEFINE_integer('image_size', 256, 'image size, default: 256') -tf.flags.DEFINE_bool('use_lsgan', True, +tf.compat.v1.flags.DEFINE_integer('batch_size', 1, 'batch size, default: 1') +tf.compat.v1.flags.DEFINE_integer('image_size', 256, 'image size, default: 256') +tf.compat.v1.flags.DEFINE_bool('use_lsgan', True, 'use lsgan (mean squared error) or cross entropy loss, default: True') -tf.flags.DEFINE_string('norm', 'instance', +tf.compat.v1.flags.DEFINE_string('norm', 'instance', '[instance, batch] use instance norm or batch norm, default: instance') -tf.flags.DEFINE_integer('lambda1', 10, +tf.compat.v1.flags.DEFINE_integer('lambda1', 10, 'weight for forward cycle loss (X->Y->X), default: 10') -tf.flags.DEFINE_integer('lambda2', 10, +tf.compat.v1.flags.DEFINE_integer('lambda2', 10, 'weight for backward cycle loss (Y->X->Y), default: 10') -tf.flags.DEFINE_float('learning_rate', 2e-4, +tf.compat.v1.flags.DEFINE_float('learning_rate', 2e-4, 'initial learning rate for Adam, default: 0.0002') -tf.flags.DEFINE_float('beta1', 0.5, +tf.compat.v1.flags.DEFINE_float('beta1', 0.5, 'momentum term of Adam, default: 0.5') -tf.flags.DEFINE_float('pool_size', 50, +tf.compat.v1.flags.DEFINE_float('pool_size', 50, 'size of image buffer that stores previously generated images, default: 50') -tf.flags.DEFINE_integer('ngf', 64, +tf.compat.v1.flags.DEFINE_integer('ngf', 64, 'number of gen filters in first conv layer, default: 64') -tf.flags.DEFINE_string('X', 'data/tfrecords/apple.tfrecords', +tf.compat.v1.flags.DEFINE_string('X', 'data/tfrecords/apple.tfrecords', 'X tfrecords file for training, default: data/tfrecords/apple.tfrecords') -tf.flags.DEFINE_string('Y', 'data/tfrecords/orange.tfrecords', +tf.compat.v1.flags.DEFINE_string('Y', 'data/tfrecords/orange.tfrecords', 'Y tfrecords file for training, default: data/tfrecords/orange.tfrecords') -tf.flags.DEFINE_string('load_model', None, - 'folder of saved model that you wish to continue training (e.g. 20170602-1936), default: None') +tf.compat.v1.flags.DEFINE_string('load_model', None, + 'folder of saved model that you wish to continue training (e.g. 20170602-1936), default: None') def train(): - if FLAGS.load_model is not None: - checkpoints_dir = "checkpoints/" + FLAGS.load_model.lstrip("checkpoints/") - else: - current_time = datetime.now().strftime("%Y%m%d-%H%M") - checkpoints_dir = "checkpoints/{}".format(current_time) - try: - os.makedirs(checkpoints_dir) - except os.error: - pass - - graph = tf.Graph() - with graph.as_default(): - cycle_gan = CycleGAN( - X_train_file=FLAGS.X, - Y_train_file=FLAGS.Y, - batch_size=FLAGS.batch_size, - image_size=FLAGS.image_size, - use_lsgan=FLAGS.use_lsgan, - norm=FLAGS.norm, - lambda1=FLAGS.lambda1, - lambda2=FLAGS.lambda2, - learning_rate=FLAGS.learning_rate, - beta1=FLAGS.beta1, - ngf=FLAGS.ngf - ) - G_loss, D_Y_loss, F_loss, D_X_loss, fake_y, fake_x = cycle_gan.model() - optimizers = cycle_gan.optimize(G_loss, D_Y_loss, F_loss, D_X_loss) - - summary_op = tf.summary.merge_all() - train_writer = tf.summary.FileWriter(checkpoints_dir, graph) - saver = tf.train.Saver() - - with tf.Session(graph=graph) as sess: if FLAGS.load_model is not None: - checkpoint = tf.train.get_checkpoint_state(checkpoints_dir) - meta_graph_path = checkpoint.model_checkpoint_path + ".meta" - restore = tf.train.import_meta_graph(meta_graph_path) - restore.restore(sess, tf.train.latest_checkpoint(checkpoints_dir)) - step = int(meta_graph_path.split("-")[2].split(".")[0]) + checkpoints_dir = "checkpoints/" + FLAGS.load_model.lstrip("checkpoints/") else: - sess.run(tf.global_variables_initializer()) - step = 0 - - coord = tf.train.Coordinator() - threads = tf.train.start_queue_runners(sess=sess, coord=coord) - - try: - fake_Y_pool = ImagePool(FLAGS.pool_size) - fake_X_pool = ImagePool(FLAGS.pool_size) - - while not coord.should_stop(): - # get previously generated images - fake_y_val, fake_x_val = sess.run([fake_y, fake_x]) - - # train - _, G_loss_val, D_Y_loss_val, F_loss_val, D_X_loss_val, summary = ( - sess.run( - [optimizers, G_loss, D_Y_loss, F_loss, D_X_loss, summary_op], - feed_dict={cycle_gan.fake_y: fake_Y_pool.query(fake_y_val), - cycle_gan.fake_x: fake_X_pool.query(fake_x_val)} - ) + current_time = datetime.now().strftime("%Y%m%d-%H%M") + checkpoints_dir = "checkpoints/{}".format(current_time) + try: + os.makedirs(checkpoints_dir) + except os.error: + pass + + graph = tf.Graph() + with graph.as_default(): + cycle_gan = CycleGAN( + X_train_file=FLAGS.X, + Y_train_file=FLAGS.Y, + batch_size=FLAGS.batch_size, + image_size=FLAGS.image_size, + use_lsgan=FLAGS.use_lsgan, + norm=FLAGS.norm, + lambda1=FLAGS.lambda1, + lambda2=FLAGS.lambda2, + learning_rate=FLAGS.learning_rate, + beta1=FLAGS.beta1, + ngf=FLAGS.ngf ) + G_loss, D_Y_loss, F_loss, D_X_loss, fake_y, fake_x = cycle_gan.model() + optimizers = cycle_gan.optimize(G_loss, D_Y_loss, F_loss, D_X_loss) + + summary_op = tf.compat.v1.summary.merge_all() + train_writer = tf.compat.v1.summary.FileWriter(checkpoints_dir, graph) + saver = tf.compat.v1.train.Saver() + + with tf.compat.v1.Session(graph=graph) as sess: + if FLAGS.load_model is not None: + checkpoint = tf.train.get_checkpoint_state(checkpoints_dir) + meta_graph_path = checkpoint.model_checkpoint_path + ".meta" + restore = tf.compat.v1.train.import_meta_graph(meta_graph_path) + restore.restore(sess, tf.train.latest_checkpoint(checkpoints_dir)) + step = int(meta_graph_path.split("-")[2].split(".")[0]) + else: + sess.run(tf.compat.v1.global_variables_initializer()) + step = 0 + + coord = tf.train.Coordinator() + threads = tf.compat.v1.train.start_queue_runners(sess=sess, coord=coord) + + try: + fake_Y_pool = ImagePool(FLAGS.pool_size) + fake_X_pool = ImagePool(FLAGS.pool_size) + + while not coord.should_stop(): + # get previously generated images + fake_y_val, fake_x_val = sess.run([fake_y, fake_x]) + + # train + _, G_loss_val, D_Y_loss_val, F_loss_val, D_X_loss_val, summary = ( + sess.run( + [optimizers, G_loss, D_Y_loss, F_loss, D_X_loss, summary_op], + feed_dict={cycle_gan.fake_y: fake_Y_pool.query(fake_y_val), + cycle_gan.fake_x: fake_X_pool.query(fake_x_val)} + ) + ) + + train_writer.add_summary(summary, step) + train_writer.flush() + + if step % 100 == 0: + logging.info('-----------Step %d:-------------' % step) + logging.info(' G_loss : {}'.format(G_loss_val)) + logging.info(' D_Y_loss : {}'.format(D_Y_loss_val)) + logging.info(' F_loss : {}'.format(F_loss_val)) + logging.info(' D_X_loss : {}'.format(D_X_loss_val)) + + if step % 10000 == 0: + save_path = saver.save(sess, checkpoints_dir + "/model.ckpt", global_step=step) + logging.info("Model saved in file: %s" % save_path) + + step += 1 + + except KeyboardInterrupt: + logging.info('Interrupted') + coord.request_stop() + except Exception as e: + coord.request_stop(e) + finally: + save_path = saver.save(sess, checkpoints_dir + "/model.ckpt", global_step=step) + logging.info("Model saved in file: %s" % save_path) + # When done, ask the threads to stop. + coord.request_stop() + coord.join(threads) - train_writer.add_summary(summary, step) - train_writer.flush() - - if step % 100 == 0: - logging.info('-----------Step %d:-------------' % step) - logging.info(' G_loss : {}'.format(G_loss_val)) - logging.info(' D_Y_loss : {}'.format(D_Y_loss_val)) - logging.info(' F_loss : {}'.format(F_loss_val)) - logging.info(' D_X_loss : {}'.format(D_X_loss_val)) - - if step % 10000 == 0: - save_path = saver.save(sess, checkpoints_dir + "/model.ckpt", global_step=step) - logging.info("Model saved in file: %s" % save_path) - - step += 1 - - except KeyboardInterrupt: - logging.info('Interrupted') - coord.request_stop() - except Exception as e: - coord.request_stop(e) - finally: - save_path = saver.save(sess, checkpoints_dir + "/model.ckpt", global_step=step) - logging.info("Model saved in file: %s" % save_path) - # When done, ask the threads to stop. - coord.request_stop() - coord.join(threads) def main(unused_argv): - train() + train() + if __name__ == '__main__': - logging.basicConfig(level=logging.INFO) - tf.app.run() + logging.basicConfig(level=logging.INFO) + tf.compat.v1.app.run() diff --git a/utils.py b/utils.py index b0aab6e..b794f0a 100644 --- a/utils.py +++ b/utils.py @@ -1,58 +1,63 @@ import tensorflow as tf import random + def convert2int(image): - """ Transfrom from float tensor ([-1.,1.]) to int image ([0,255]) + """ Transfrom from float tensor ([-1.,1.]) to int image ([0,255]) """ - return tf.image.convert_image_dtype((image+1.0)/2.0, tf.uint8) + return tf.image.convert_image_dtype((image + 1.0) / 2.0, tf.uint8) + def convert2float(image): - """ Transfrom from int image ([0,255]) to float tensor ([-1.,1.]) + """ Transfrom from int image ([0,255]) to float tensor ([-1.,1.]) """ - image = tf.image.convert_image_dtype(image, dtype=tf.float32) - return (image/127.5) - 1.0 + image = tf.image.convert_image_dtype(image, dtype=tf.float32) + return (image / 127.5) - 1.0 + def batch_convert2int(images): - """ + """ Args: images: 4D float tensor (batch_size, image_size, image_size, depth) Returns: 4D int tensor """ - return tf.map_fn(convert2int, images, dtype=tf.uint8) + return tf.map_fn(convert2int, images, dtype=tf.uint8) + def batch_convert2float(images): - """ + """ Args: images: 4D int tensor (batch_size, image_size, image_size, depth) Returns: 4D float tensor """ - return tf.map_fn(convert2float, images, dtype=tf.float32) + return tf.map_fn(convert2float, images, dtype=tf.float32) + class ImagePool: - """ History of generated images + """ History of generated images Same logic as https://github.com/junyanz/CycleGAN/blob/master/util/image_pool.lua """ - def __init__(self, pool_size): - self.pool_size = pool_size - self.images = [] - - def query(self, image): - if self.pool_size == 0: - return image - - if len(self.images) < self.pool_size: - self.images.append(image) - return image - else: - p = random.random() - if p > 0.5: - # use old image - random_id = random.randrange(0, self.pool_size) - tmp = self.images[random_id].copy() - self.images[random_id] = image.copy() - return tmp - else: - return image + def __init__(self, pool_size): + self.pool_size = pool_size + self.images = [] + + def query(self, image): + if self.pool_size == 0: + return image + + if len(self.images) < self.pool_size: + self.images.append(image) + return image + else: + p = random.random() + if p > 0.5: + # use old image + random_id = random.randrange(0, self.pool_size) + tmp = self.images[random_id].copy() + self.images[random_id] = image.copy() + return tmp + else: + return image