Add Dataset class

src/data/dataset.py

@@ -3,101 +3,106 @@
 from __future__ import print_function
 
 import os
-import multiprocessing
 
 import tensorflow as tf
 
-def load_image(image_file, img_size=256):
- # Read file into tensor of type string.
- image_string = tf.read_file(image_file)
- # Decodes file into jpg of type uint8 (range [0, 255]).
- image = tf.image.decode_jpeg(image_string, channels=3)
- # Convert to floating point with 32 bits (range [0, 1]).
- image = tf.image.convert_image_dtype(image, tf.float32)
- # Resize with bicubic interpolation, making sure that corner pixel values
- # are preserved.
- image = tf.image.resize_images(image, size=[img_size, img_size],
- method=tf.image.ResizeMethod.BICUBIC, align_corners=True)
- # Transform image to [-1, 1] from [0, 1].
- image = (image - 0.5) * 2
- return image
+def Dataset(object):
 
-def save_images(image_to_save, save_dir, image_index, img_size=256):
- save_file = os.path.join(save_dir,'test' + str(image_index) + '.jpg')
- # Reshape to get rid of batch size dimension in the tensor.
- image = tf.reshape(image_to_save, shape=[img_size, img_size, 3])
- # Scale from [-1, 1] to [0, 1).
- image = (image * 0.5) + 0.5
- # Convert to uint8 (range [0, 255]), saturate to avoid possible under/overflow.
- image = tf.image.convert_image_dtype(image, dtype=tf.uint8, saturate=True)
- # JPEG encode image into string Tensor.
- image_string = tf.image.encode_jpeg(image, format='rgb', quality=95)
- tf.write_file(filename=save_file, contents=image_string)
+ def __init__(self, opt):
+ self.opt = opt
+ self.gpu_id = "/gpu:" + str(self.opt.gpu_id)
+ if opt.training:
+ self.trainA_path = os.path.join(self.opt.data_dir, 'trainA')
+ self.trainB_path = os.path.join(self.opt.data_dir, 'trainB')
+ self.trainA_size = len(os.listdir(self.trainA_path))
+ self.trainB_size = len(os.listdir(self.trainB_path))
+ dataA, dataB = self.load_train_data()
+ else:
+ self.testA_path = os.path.join(self.opt.data_dir, 'testA')
+ self.testB_path = os.path.join(self.opt.data_dir, 'testB')
+ self.testA_size = len(os.listdir(self.testA_path))
+ self.testB_size = len(os.listdir(self.testB_path))
+ dataA, dataB = self.load_test_data()
+ self.data = {"A": dataA, "B": dataB}
 
-# TODO: Refactor into class
-def load_train_data(dataset_id, project_dir, batch_size=1):
- path_to_dataset = os.path.join(project_dir, 'data', 'raw', dataset_id + os.sep)
- trainA_path = os.path.join(path_to_dataset, 'trainA')
- trainB_path = os.path.join(path_to_dataset, 'trainB')
- trainA_size = len(os.listdir(trainA_path))
- trainB_size = len(os.listdir(trainB_path))
- threads = multiprocessing.cpu_count()
+ def load_train_data(self):
+ # Create Dataset from folder of string filenames.
+ train_datasetA = tf.data.Dataset.list_files(self.trainA_path + os.sep + '*.jpg', shuffle=False)
+ train_datasetB = tf.data.Dataset.list_files(self.trainB_path + os.sep + '*.jpg', shuffle=False)
+ # Infinitely loop the dataset, shuffling once per epoch (in memory).
+ # Safe to do since the dataset pipeline is currently string filenames.
+ # Fused operation is faster than separated shuffle and repeat.
+ train_datasetA = train_datasetA.apply(tf.contrib.data.shuffle_and_repeat(buffer_size=self.trainA_size))
+ train_datasetB = train_datasetB.apply(tf.contrib.data.shuffle_and_repeat(buffer_size=self.trainB_size))
+ # Decodes filenames into jpegs, then stacks them into batches.
+ # Throwing away the remainder allows the pipeline to report a fixed sized
+ # batch size, aiding in model definition downstream.
+ train_datasetA = train_datasetA.apply(tf.contrib.data.map_and_batch(lambda x: self.load_image(x),
+ batch_size=self.opt.batch_size,
+ num_parallel_calls=self.opt.num_threads,
+ drop_remainder=True))
+ train_datasetB = train_datasetB.apply(tf.contrib.data.map_and_batch(lambda x: self.load_image(x),
+ batch_size=self.opt.batch_size,
+ num_parallel_calls=self.opt.num_threads,
+ drop_remainder=True))
+ # Queue up a number of batches on CPU side:
+ train_datasetA = train_datasetA.prefetch(buffer_size=self.opt.num_threads)
+ train_datasetB = train_datasetB.prefetch(buffer_size=self.opt.num_threads)
+ # Queue up batches asynchronously onto the GPU.
+ # As long as there is a pool of batches CPU side a GPU prefetch of 1 is fine.
+ # If no GPU exists gpu_id = -1:
+ if self.opt.gpu_id != -1:
+ train_datasetA = train_datasetA.apply(tf.contrib.data.prefetch_to_device(self.gpu_id, buffer_size=1))
+ train_datasetB = train_datasetB.apply(tf.contrib.data.prefetch_to_device(self.gpu_id, buffer_size=1))
+ # Create a tf.data.Iterator from the Datasets:
+ return iter(train_datasetA), iter(train_datasetB)
 
- # Create Dataset from folder of string filenames.
- train_datasetA = tf.data.Dataset.list_files(trainA_path + os.sep + '*.jpg', shuffle=False)
- # Infinitely loop the dataset, shuffling once per epoch (in memory).
- # Safe to do since the dataset pipeline is currently string filenames.
- # Fused operation is faster than separated shuffle and repeat.
- # This is also serializable, so Dataset state can be saved with Checkpoints,
- # but doing this causes a segmentation fault for some reason...
- train_datasetA = train_datasetA.apply(tf.contrib.data.shuffle_and_repeat(buffer_size=trainA_size))
- # Decodes filenames into jpegs, then stacks them into batches.
- # Throwing away the remainder allows the pipeline to report a fixed sized
- # batch size, aiding in model definition downstream.
- train_datasetA = train_datasetA.apply(tf.contrib.data.map_and_batch(lambda x: load_image(x),
- batch_size=batch_size,
- num_parallel_calls=threads,
- drop_remainder=True))
- # Queue up a number of batches on CPU side
- train_datasetA = train_datasetA.prefetch(buffer_size=threads)
- # Queue up batches asynchronously onto the GPU.
- # As long as there is a pool of batches CPU side a GPU prefetch of 1 is fine.
- # TODO: If GPU exists:
- train_datasetA = train_datasetA.apply(tf.contrib.data.prefetch_to_device("/gpu:0", buffer_size=1))
+ def load_test_data(self):
+ test_datasetA = tf.data.Dataset.list_files(self.testA_path + os.sep + '*.jpg', shuffle=False)
+ test_datasetB = tf.data.Dataset.list_files(self.testB_path + os.sep + '*.jpg', shuffle=False)
+ test_datasetA = test_datasetA.apply(tf.contrib.data.map_and_batch(lambda x: self.load_image(x),
+ batch_size=1,
+ num_parallel_calls=self.opt.num_threads,
+ drop_remainder=False))
+ test_datasetB = test_datasetB.apply(tf.contrib.data.map_and_batch(lambda x: self.load_image(x),
+ batch_size=1,
+ num_parallel_calls=self.opt.num_threads,
+ drop_remainder=False))
+ test_datasetA = test_datasetA.prefetch(buffer_size=self.opt.num_threads)
+ test_datasetB = test_datasetB.prefetch(buffer_size=self.opt.num_threads)
+ if self.opt.gpu_id != -1:
+ train_datasetA = train_datasetA.apply(tf.contrib.data.prefetch_to_device(self.gpu_id, buffer_size=1))
+ train_datasetB = train_datasetB.apply(tf.contrib.data.prefetch_to_device(self.gpu_id, buffer_size=1))
+ return iter(test_datasetA), iter(test_datasetB)
 
- train_datasetB = tf.data.Dataset.list_files(trainB_path + os.sep + '*.jpg', shuffle=False)
- train_datasetB = train_datasetB.apply(tf.contrib.data.shuffle_and_repeat(buffer_size=trainB_size))
- train_datasetB = train_datasetB.apply(tf.contrib.data.map_and_batch(lambda x: load_image(x),
- batch_size=batch_size,
- num_parallel_calls=threads,
- drop_remainder=True))
- train_datasetB = train_datasetB.prefetch(buffer_size=threads)
- train_datasetB = train_datasetB.apply(tf.contrib.data.prefetch_to_device("/gpu:0", buffer_size=1))
- # Create a tf.data.Iterator from the Datasets:
- return iter(train_datasetA), iter(train_datasetB)
+ def load_image(self, image_file):
+ # Read file into tensor of type string.
+ image_string = tf.read_file(image_file)
+ # Decodes file into jpg of type uint8 (range [0, 255]).
+ image = tf.image.decode_jpeg(image_string, channels=3)
+ # Convert to floating point with 32 bits (range [0, 1]).
+ image = tf.image.convert_image_dtype(image, tf.float32)
+ # Resize with bicubic interpolation, making sure that corner pixel values
+ # are preserved.
+ image = tf.image.resize_images(image, size=[self.opt.img_size, self.opt.img_size],
+ method=tf.image.ResizeMethod.BICUBIC, align_corners=True)
+ # Transform image to [-1, 1] from [0, 1].
+ image = (image - 0.5) * 2
+ return image
 
-def load_test_data(dataset_id, project_dir):
- path_to_dataset = os.path.join(project_dir, 'data', 'raw', dataset_id + os.sep)
- testA_path = os.path.join(path_to_dataset, 'testA')
- testB_path = os.path.join(path_to_dataset, 'testB')
- testA_size = len(os.listdir(testA_path))
- testB_size = len(os.listdir(testB_path))
- threads = multiprocessing.cpu_count()
+ def save_images(self, image_to_save, save_dir, image_index):
+ save_file = os.path.join(save_dir,'test' + str(image_index) + '.jpg')
+ # Reshape to get rid of batch size dimension in the tensor.
+ image = tf.reshape(image_to_save, shape=[self.opt.img_size, self.opt.img_size, 3])
+ # Scale from [-1, 1] to [0, 1).
+ image = (image * 0.5) + 0.5
+ # Convert to uint8 (range [0, 255]), saturate to avoid possible under/overflow.
+ image = tf.image.convert_image_dtype(image, dtype=tf.uint8, saturate=True)
+ # JPEG encode image into string Tensor.
+ image_string = tf.image.encode_jpeg(image, format='rgb', quality=95)
+ tf.write_file(filename=save_file, contents=image_string)
 
- test_datasetA = tf.data.Dataset.list_files(testA_path + os.sep + '*.jpg', shuffle=False)
- test_datasetA = test_datasetA.apply(tf.contrib.data.map_and_batch(lambda x: load_image(x),
- batch_size=1,
- num_parallel_calls=threads,
- drop_remainder=False))
- test_datasetA = test_datasetA.prefetch(buffer_size=threads)
- test_datasetA = test_datasetA.apply(tf.contrib.data.prefetch_to_device("/gpu:0", buffer_size=1))
-
- test_datasetB = tf.data.Dataset.list_files(testB_path + os.sep + '*.jpg', shuffle=False)
- test_datasetB = test_datasetB.apply(tf.contrib.data.map_and_batch(lambda x: load_image(x),
- batch_size=1,
- num_parallel_calls=threads,
- drop_remainder=False))
- test_datasetB = test_datasetB.prefetch(buffer_size=threads)
- test_datasetB = test_datasetB.apply(tf.contrib.data.prefetch_to_device("/gpu:0", buffer_size=1))
-
- return test_datasetA, test_datasetB, testA_size, testB_size
+ def get_batches_per_epoch(self, opt):
+ # floor(Avg dataset size / batch_size)
+ batches_per_epoch = (self.trainA_size + self.trainB_size) // (2 * opt.batch_size)
+ return batches_per_epoch

src/models/__init__.py

@@ -1,14 +1 @@
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
 
-import os
-
-import tensorflow as tf
-
-# TODO: Merge into dataset class
-def get_batches_per_epoch(opt):
- trainA_size = len(os.listdir(os.path.join(opt.data_dir, 'trainA')))
- trainB_size = len(os.listdir(os.path.join(opt.data_dir, 'trainB')))
- batches_per_epoch = (trainA_size + trainB_size) // (2 * opt.batch_size) # floor(Avg dataset size / batch_size)
- return batches_per_epoch

src/models/cyclegan.py

@@ -62,8 +62,9 @@ def restore_checkpoint(self):
  print("No checkpoint found, initializing model.")
 
  def set_input(self, input):
- self.dataA = input["A"]
- self.dataB = input["B"]
+ # Get next batches:
+ self.dataA = input["A"].get_next()
+ self.dataB = input["B"].get_next()
 
  def forward(self):
  # Gen output shape: (batch_size, img_size, img_size, 3)

src/train.py

@@ -7,37 +7,32 @@
 
 import tensorflow as tf
 
-import models
 from models.cyclegan import CycleGANModel
-from data.dataset import load_train_data
+from data.dataset import Dataset
 from utils.options import Options
 
 tf.enable_eager_execution()
 
-def train_one_epoch(self):
- raise NotImplementedError
-
 if __name__ == "__main__":
  opt = Options(isTrain=True)
- with tf.device("/cpu:0"):
- # Preprocess data on CPU for significant performance gains:
- dataA, dataB = load_train_data(dataset_id, project_dir)
- model = CycleGANModel(opt)
- batches_per_epoch = models.get_batches_per_epoch(opt)
- with tf.device("/gpu:0"):
+ # TODO: Test if this is always on CPU:
+ dataset = Dataset(opt)
+ model = CycleGANModel(opt)
+ # TODO: Figure out GPU conditionals:
+ with tf.device("/gpu:" + str(opt.gpu_id)):
  global_step = model.global_step
+ batches_per_epoch = dataset.get_batches_per_epoch(opt)
  # Initialize Tensorboard summary writer:
- log_dir = os.path.join(opt.save_dir, 'tensorboard')
+ log_dir = os.path.join(opt.save_dir, 'tensorboard') # TODO: move to model class?
  summary_writer = tf.contrib.summary.create_file_writer(log_dir, flush_millis=10000)
- for epoch in range(opt.epochs):
+ for epoch in range(1, opt.epochs):
  start = time.time()
  with summary_writer.as_default():
  for train_step in range(batches_per_epoch):
  # Record summaries every 100 train_steps, we multiply by 3 because there are 3 gradient updates per step.
  with tf.contrib.summary.record_summaries_every_n_global_steps(opt.summary_freq * 3, global_step=global_step):
  # Get next training batches:
- batch = {"A": dataA.get_next(), "B": dataB.get_next()}
- model.set_input(batch)
+ model.set_input(dataset.data)
  model.optimize_parameters()
 
  # Summaries for Tensorboard:
@@ -55,7 +50,7 @@ def train_one_epoch(self):
  # Assign decayed learning rate:
  model.update_learning_rate(batches_per_epoch)
  # Checkpoint the model:
- if (epoch + 1) % opt.save_epoch_freq == 0:
+ if epoch % opt.save_epoch_freq == 0:
  model.save_model()
  print("Global Training Step: ", global_step.numpy() // 3)
  print ("Time taken for total epoch {} is {} sec\n".format(global_step.numpy() \