From 2f16b7d724fae60ff75593ada5616fad0a029e16 Mon Sep 17 00:00:00 2001 From: Padmaja Bhagwat Date: Sat, 24 Jul 2021 14:43:19 +0530 Subject: [PATCH 1/3] making the code tensorflow v2.5 compatible --- build_data.py | 124 ++++++++++---------- discriminator.py | 51 +++++---- export_graph.py | 51 +++++---- generator.py | 95 ++++++++-------- model.py | 291 ++++++++++++++++++++++++----------------------- ops.py | 258 +++++++++++++++++++++-------------------- reader.py | 137 +++++++++++----------- train.py | 182 ++++++++++++++--------------- utils.py | 67 ++++++----- 9 files changed, 644 insertions(+), 612 deletions(-) diff --git a/build_data.py b/build_data.py index fb4f2eb6..217ae934 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 162ec6d4..4ddb92ea 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.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.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope=self.name) - return output + return output diff --git a/export_graph.py b/export_graph.py index 7f94a6f2..2cc3b5c9 100644 --- a/export_graph.py +++ b/export_graph.py @@ -24,37 +24,40 @@ tf.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.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.train.Saver() + export_saver = tf.train.Saver() - output_image = tf.identity(output_image, name='output_image') - restore_saver = tf.train.Saver() - export_saver = tf.train.Saver() + 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]) - 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.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.app.run() diff --git a/generator.py b/generator.py index ecb715b5..ffd61360 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.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 diff --git a/model.py b/model.py index 98a9af19..232f03e3 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.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) 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.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), 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.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.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 9becd248..bc7e2568 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.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.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.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 + 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.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.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.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.get_variable( + name, shape, + initializer=tf.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.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.variable_scope("batch_norm"): + return tf.contrib.layers.batch_norm(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.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 + def safe_log(x, eps=1e-12): - return tf.log(x + eps) + return tf.log(x + eps) diff --git a/reader.py b/reader.py index f61e2851..06911399 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.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.train.string_input_producer([self.tfrecords_file]) + reader = tf.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.parse_single_example( + serialized_example, + features={ + 'image/file_name': tf.FixedLenFeature([], tf.string), + 'image/encoded_image': tf.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.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_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 _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 0b0cc272..f1e5cd20 100644 --- a/train.py +++ b/train.py @@ -32,104 +32,106 @@ tf.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') + '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.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]) + 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)} + ) + ) + + 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.app.run() diff --git a/utils.py b/utils.py index b0aab6e7..b794f0a1 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 From 6a7e050bcd8fac18658a8c973c39aa318fdefad9 Mon Sep 17 00:00:00 2001 From: Padmaja Bhagwat Date: Sat, 24 Jul 2021 16:21:35 +0530 Subject: [PATCH 2/3] making the code tensorflow v2.5 compatible --- build_data.py | 16 ++++++------- discriminator.py | 4 ++-- export_graph.py | 34 +++++++++++++-------------- generator.py | 4 ++-- inference.py | 61 +++++++++++++++++++++++++----------------------- model.py | 14 +++++------ ops.py | 44 +++++++++++++++++----------------- reader.py | 20 ++++++++-------- train.py | 49 +++++++++++++++++++------------------- 9 files changed, 124 insertions(+), 122 deletions(-) diff --git a/build_data.py b/build_data.py index 217ae934..7f0684cd 100644 --- a/build_data.py +++ b/build_data.py @@ -10,14 +10,14 @@ FLAGS = tf.compat.v1.flags.FLAGS -tf.compat.v1.flags.DEFINE_string('X_input_dir', 'data/apple2orange/trainA', - 'X input directory, default: data/apple2orange/trainA') -tf.compat.v1.flags.DEFINE_string('Y_input_dir', 'data/apple2orange/trainB', - 'Y input directory, default: data/apple2orange/trainB') -tf.compat.v1.flags.DEFINE_string('X_output_file', 'data/tfrecords/apple.tfrecords', - 'X output tfrecords file, default: data/tfrecords/apple.tfrecords') -tf.compat.v1.flags.DEFINE_string('Y_output_file', 'data/tfrecords/orange.tfrecords', - 'Y output tfrecords file, default: data/tfrecords/orange.tfrecords') +tf.compat.v1.flags.DEFINE_string('X_input_dir', 'data/ukiyoe2photo/trainA', + 'X input directory, default: data/ukiyoe2photo/trainA') +tf.compat.v1.flags.DEFINE_string('Y_input_dir', 'data/ukiyoe2photo/trainB', + 'Y input directory, default: data/ukiyoe2photo/trainB') +tf.compat.v1.flags.DEFINE_string('X_output_file', 'data/tfrecords/ukiyoe.tfrecords', + 'X output tfrecords file, default: data/tfrecords/ukiyoe.tfrecords') +tf.compat.v1.flags.DEFINE_string('Y_output_file', 'data/tfrecords/photo.tfrecords', + 'Y output tfrecords file, default: data/tfrecords/photo.tfrecords') def data_reader(input_dir, shuffle=True): diff --git a/discriminator.py b/discriminator.py index 4ddb92ea..9f7b2616 100644 --- a/discriminator.py +++ b/discriminator.py @@ -18,7 +18,7 @@ def __call__(self, input): 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): + 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) @@ -35,6 +35,6 @@ def __call__(self, input): 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.variables = tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.TRAINABLE_VARIABLES, scope=self.name) return output diff --git a/export_graph.py b/export_graph.py index 2cc3b5c9..4ce0cc5b 100644 --- a/export_graph.py +++ b/export_graph.py @@ -2,8 +2,8 @@ This makes file size smaller and can be used for inference in production. An example of command-line usage is: python export_graph.py --checkpoint_dir checkpoints/20170424-1152 \ - --XtoY_model apple2orange.pb \ - --YtoX_model orange2apple.pb \ + --XtoY_model ukiyoe2photo.pb \ + --YtoX_model photo2ukiyoe.pb \ --image_size 256 """ @@ -13,15 +13,15 @@ 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', 'ukiyoe2photo.pb', 'XtoY model name, default: ukiyoe2photo.pb') +tf.compat.v1.flags.DEFINE_string('YtoX_model', 'photo2ukiyoe.pb', 'YtoX model name, default: photo2ukiyoe.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') @@ -31,7 +31,7 @@ def export_graph(model_name, XtoY=True): with graph.as_default(): cycle_gan = CycleGAN(ngf=FLAGS.ngf, norm=FLAGS.norm, image_size=FLAGS.image_size) - input_image = tf.placeholder(tf.float32, shape=[FLAGS.image_size, FLAGS.image_size, 3], name='input_image') + 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)) @@ -39,17 +39,17 @@ def export_graph(model_name, XtoY=True): output_image = cycle_gan.F.sample(tf.expand_dims(input_image, 0)) output_image = tf.identity(output_image, name='output_image') - restore_saver = tf.train.Saver() - export_saver = tf.train.Saver() + restore_saver = tf.compat.v1.train.Saver() + export_saver = tf.compat.v1.train.Saver() - with tf.Session(graph=graph) as sess: - sess.run(tf.global_variables_initializer()) + 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.graph_util.convert_variables_to_constants( + output_graph_def = tf.compat.v1.graph_util.convert_variables_to_constants( sess, graph.as_graph_def(), [output_image.op.name]) - tf.train.write_graph(output_graph_def, 'pretrained', model_name, as_text=False) + tf.compat.v1.train.write_graph(output_graph_def, 'pretrained', model_name, as_text=False) def main(unused_argv): @@ -60,4 +60,4 @@ def main(unused_argv): if __name__ == '__main__': - tf.app.run() + tf.compat.v1.app.run() diff --git a/generator.py b/generator.py index ffd61360..5e2e510d 100644 --- a/generator.py +++ b/generator.py @@ -19,7 +19,7 @@ def __call__(self, input): Returns: output: same size as input """ - with tf.variable_scope(self.name): + 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) @@ -48,7 +48,7 @@ def __call__(self, input): 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) + self.variables = tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.TRAINABLE_VARIABLES, scope=self.name) return output diff --git a/inference.py b/inference.py index 490c763e..fd1ee9ff 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 232f03e3..cd755716 100644 --- a/model.py +++ b/model.py @@ -47,7 +47,7 @@ def __init__(self, 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.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', @@ -56,9 +56,9 @@ def __init__(self, self.D_X = Discriminator('D_X', self.is_training, norm=norm, use_sigmoid=use_sigmoid) - self.fake_x = tf.placeholder(tf.float32, + self.fake_x = tf.compat.v1.placeholder(tf.float32, shape=[batch_size, image_size, image_size, 3]) - self.fake_y = tf.placeholder(tf.float32, + self.fake_y = tf.compat.v1.placeholder(tf.float32, shape=[batch_size, image_size, image_size, 3]) def model(self): @@ -117,7 +117,7 @@ def make_optimizer(loss, variables, name='Adam'): learning_rate = ( tf.where( tf.greater_equal(global_step, start_decay_step), - tf.train.polynomial_decay(starter_learning_rate, 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 @@ -127,7 +127,7 @@ def make_optimizer(loss, variables, name='Adam'): tf.summary.scalar('learning_rate/{}'.format(name), learning_rate) learning_step = ( - tf.train.AdamOptimizer(learning_rate, beta1=beta1, name=name) + tf.compat.v1.train.AdamOptimizer(learning_rate, beta1=beta1, name=name) .minimize(loss, global_step=global_step, var_list=variables) ) return learning_step @@ -152,7 +152,7 @@ def discriminator_loss(self, D, y, fake_y, use_lsgan=True): """ if use_lsgan: # use mean squared error - error_real = tf.reduce_mean(tf.squared_difference(D(y), REAL_LABEL)) + 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 @@ -166,7 +166,7 @@ def generator_loss(self, D, fake_y, use_lsgan=True): """ if use_lsgan: # use mean squared error - loss = tf.reduce_mean(tf.squared_difference(D(fake_y), REAL_LABEL)) + 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 diff --git a/ops.py b/ops.py index bc7e2568..79f87f0e 100644 --- a/ops.py +++ b/ops.py @@ -17,7 +17,7 @@ def c7s1_k(input, k, reuse=False, norm='instance', activation='relu', is_trainin Returns: 4D tensor """ - with tf.variable_scope(name, reuse=reuse): + with tf.compat.v1.variable_scope(name, reuse=reuse): weights = _weights("weights", shape=[7, 7, input.get_shape()[3], k]) @@ -47,7 +47,7 @@ def dk(input, k, reuse=False, norm='instance', is_training=True, name=None): Returns: 4D tensor """ - with tf.variable_scope(name, reuse=reuse): + with tf.compat.v1.variable_scope(name, reuse=reuse): weights = _weights("weights", shape=[3, 3, input.get_shape()[3], k]) @@ -69,8 +69,8 @@ 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): + 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') @@ -79,7 +79,7 @@ def Rk(input, k, reuse=False, norm='instance', is_training=True, name=None): normalized1 = _norm(conv1, is_training, norm) relu1 = tf.nn.relu(normalized1) - with tf.variable_scope('layer2', reuse=reuse): + with tf.compat.v1.variable_scope('layer2', reuse=reuse): weights2 = _weights("weights2", shape=[3, 3, relu1.get_shape()[3], k]) @@ -113,7 +113,7 @@ def uk(input, k, reuse=False, norm='instance', is_training=True, name=None, outp Returns: 4D tensor """ - with tf.variable_scope(name, reuse=reuse): + with tf.compat.v1.variable_scope(name, reuse=reuse): input_shape = input.get_shape().as_list() weights = _weights("weights", @@ -145,7 +145,7 @@ 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): + with tf.compat.v1.variable_scope(name, reuse=reuse): weights = _weights("weights", shape=[4, 4, input.get_shape()[3], k]) @@ -166,7 +166,7 @@ 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): + with tf.compat.v1.variable_scope(name, reuse=reuse): weights = _weights("weights", shape=[4, 4, input.get_shape()[3], 1]) biases = _biases("biases", [1]) @@ -190,9 +190,9 @@ def _weights(name, shape, mean=0.0, stddev=0.02): Returns: A trainable variable """ - var = tf.get_variable( + var = tf.compat.v1.get_variable( name, shape, - initializer=tf.random_normal_initializer( + initializer=tf.compat.v1.random_normal_initializer( mean=mean, stddev=stddev, dtype=tf.float32)) return var @@ -200,8 +200,8 @@ def _weights(name, shape, mean=0.0, stddev=0.02): def _biases(name, shape, constant=0.0): """ 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): @@ -222,27 +222,27 @@ def _norm(input, is_training, norm='instance'): def _batch_norm(input, is_training): """ 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 """ - with tf.variable_scope("instance_norm"): + 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], keep_dims=True) + mean, variance = tf.nn.moments(input, axes=[1, 2], keepdims=True) epsilon = 1e-5 - inv = tf.rsqrt(variance + epsilon) + 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 06911399..bec17c50 100644 --- a/reader.py +++ b/reader.py @@ -17,7 +17,7 @@ def __init__(self, tfrecords_file, image_size=256, self.min_queue_examples = min_queue_examples self.batch_size = batch_size self.num_threads = num_threads - self.reader = tf.TFRecordReader() + self.reader = tf.compat.v1.TFRecordReader() self.name = name def feed(self): @@ -26,21 +26,21 @@ def feed(self): 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() + 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( + features = tf.compat.v1.parse_single_example( serialized_example, features={ - 'image/file_name': tf.FixedLenFeature([], tf.string), - 'image/encoded_image': tf.FixedLenFeature([], tf.string), + '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( + 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 @@ -50,15 +50,15 @@ def feed(self): return images def _preprocess(self, image): - image = tf.image.resize_images(image, size=(self.image_size, self.image_size)) + 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 test_reader(): - TRAIN_FILE_1 = 'data/tfrecords/apple.tfrecords' - TRAIN_FILE_2 = 'data/tfrecords/orange.tfrecords' + TRAIN_FILE_1 = 'data/tfrecords/ukiyoe.tfrecords' + TRAIN_FILE_2 = 'data/tfrecords/photo.tfrecords' with tf.Graph().as_default(): reader1 = Reader(TRAIN_FILE_1, batch_size=2) diff --git a/train.py b/train.py index f1e5cd20..e2db5ac7 100644 --- a/train.py +++ b/train.py @@ -1,37 +1,36 @@ 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', - 'X tfrecords file for training, default: data/tfrecords/apple.tfrecords') -tf.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, +tf.compat.v1.flags.DEFINE_string('X', 'data/tfrecords/ukiyoe.tfrecords', + 'X tfrecords file for training, default: data/tfrecords/ukiyoe.tfrecords') +tf.compat.v1.flags.DEFINE_string('Y', 'data/tfrecords/photo.tfrecords', + 'Y tfrecords file for training, default: data/tfrecords/photo.tfrecords') +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') @@ -64,23 +63,23 @@ def train(): 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() + 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.Session(graph=graph) as sess: + 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.train.import_meta_graph(meta_graph_path) + 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.global_variables_initializer()) + sess.run(tf.compat.v1.global_variables_initializer()) step = 0 coord = tf.train.Coordinator() - threads = tf.train.start_queue_runners(sess=sess, coord=coord) + threads = tf.compat.v1.train.start_queue_runners(sess=sess, coord=coord) try: fake_Y_pool = ImagePool(FLAGS.pool_size) @@ -134,4 +133,4 @@ def main(unused_argv): if __name__ == '__main__': logging.basicConfig(level=logging.INFO) - tf.app.run() + tf.compat.v1.app.run() From d42eedd72f3b3e043e710ba845ef9f4912ec18d0 Mon Sep 17 00:00:00 2001 From: Padmaja Bhagwat Date: Thu, 29 Jul 2021 23:18:21 +0530 Subject: [PATCH 3/3] making the code Tensorflow 2.5 compatible --- build_data.py | 16 ++++++++-------- export_graph.py | 8 ++++---- reader.py | 4 ++-- train.py | 8 ++++---- 4 files changed, 18 insertions(+), 18 deletions(-) diff --git a/build_data.py b/build_data.py index 7f0684cd..217ae934 100644 --- a/build_data.py +++ b/build_data.py @@ -10,14 +10,14 @@ FLAGS = tf.compat.v1.flags.FLAGS -tf.compat.v1.flags.DEFINE_string('X_input_dir', 'data/ukiyoe2photo/trainA', - 'X input directory, default: data/ukiyoe2photo/trainA') -tf.compat.v1.flags.DEFINE_string('Y_input_dir', 'data/ukiyoe2photo/trainB', - 'Y input directory, default: data/ukiyoe2photo/trainB') -tf.compat.v1.flags.DEFINE_string('X_output_file', 'data/tfrecords/ukiyoe.tfrecords', - 'X output tfrecords file, default: data/tfrecords/ukiyoe.tfrecords') -tf.compat.v1.flags.DEFINE_string('Y_output_file', 'data/tfrecords/photo.tfrecords', - 'Y output tfrecords file, default: data/tfrecords/photo.tfrecords') +tf.compat.v1.flags.DEFINE_string('X_input_dir', 'data/apple2orange/trainA', + 'X input directory, default: data/apple2orange/trainA') +tf.compat.v1.flags.DEFINE_string('Y_input_dir', 'data/apple2orange/trainB', + 'Y input directory, default: data/apple2orange/trainB') +tf.compat.v1.flags.DEFINE_string('X_output_file', 'data/tfrecords/apple.tfrecords', + 'X output tfrecords file, default: data/tfrecords/apple.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): diff --git a/export_graph.py b/export_graph.py index 4ce0cc5b..c7bc2da3 100644 --- a/export_graph.py +++ b/export_graph.py @@ -2,8 +2,8 @@ This makes file size smaller and can be used for inference in production. An example of command-line usage is: python export_graph.py --checkpoint_dir checkpoints/20170424-1152 \ - --XtoY_model ukiyoe2photo.pb \ - --YtoX_model photo2ukiyoe.pb \ + --XtoY_model apple2orange.pb \ + --YtoX_model orange2apple.pb \ --image_size 256 """ @@ -16,8 +16,8 @@ FLAGS = tf.compat.v1.flags.FLAGS tf.compat.v1.flags.DEFINE_string('checkpoint_dir', '', 'checkpoints directory path') -tf.compat.v1.flags.DEFINE_string('XtoY_model', 'ukiyoe2photo.pb', 'XtoY model name, default: ukiyoe2photo.pb') -tf.compat.v1.flags.DEFINE_string('YtoX_model', 'photo2ukiyoe.pb', 'YtoX model name, default: photo2ukiyoe.pb') +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') diff --git a/reader.py b/reader.py index bec17c50..ce822e20 100644 --- a/reader.py +++ b/reader.py @@ -57,8 +57,8 @@ def _preprocess(self, image): def test_reader(): - TRAIN_FILE_1 = 'data/tfrecords/ukiyoe.tfrecords' - TRAIN_FILE_2 = 'data/tfrecords/photo.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) diff --git a/train.py b/train.py index e2db5ac7..a042564a 100644 --- a/train.py +++ b/train.py @@ -26,10 +26,10 @@ tf.compat.v1.flags.DEFINE_integer('ngf', 64, 'number of gen filters in first conv layer, default: 64') -tf.compat.v1.flags.DEFINE_string('X', 'data/tfrecords/ukiyoe.tfrecords', - 'X tfrecords file for training, default: data/tfrecords/ukiyoe.tfrecords') -tf.compat.v1.flags.DEFINE_string('Y', 'data/tfrecords/photo.tfrecords', - 'Y tfrecords file for training, default: data/tfrecords/photo.tfrecords') +tf.compat.v1.flags.DEFINE_string('X', 'data/tfrecords/apple.tfrecords', + 'X tfrecords file for training, default: data/tfrecords/apple.tfrecords') +tf.compat.v1.flags.DEFINE_string('Y', 'data/tfrecords/orange.tfrecords', + 'Y tfrecords file for training, default: data/tfrecords/orange.tfrecords') 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')