print 'NUM_CRITIC: ', NUM_CRITIC print 'LOAD_MODEL: ', LOAD_MODEL print 'JITTER: ', JITTER print 'SIZE: ', SIZE print 'L1_WEIGHT: ', L1_WEIGHT print 'L2_WEIGHT: ', L2_WEIGHT print 'GAN_WEIGHT: ', GAN_WEIGHT print # global step that is saved with a model to keep track of how many steps/epochs global_step = tf.Variable(0, name='global_step', trainable=False) # load data Data = data_ops.loadData(DATA_DIR, DATASET, BATCH_SIZE, jitter=JITTER, SIZE=SIZE) # number of training images num_train = Data.count # The gray 'lightness' channel in range [-1, 1] L_image = Data.inputs # The color channels in [-1, 1] range ab_image = Data.targets lab_images = tf.concat([L_image, ab_image], axis=3) if ARCHITECTURE == 'pix2pix': # generated ab values from generator
ckpt = tf.train.get_checkpoint_state(experiment_dir) if ckpt and ckpt.model_checkpoint_path: print "Restoring previous model..." try: saver.restore(sess, ckpt.model_checkpoint_path) print "Model restored" except: print "Could not restore model" pass step = int(sess.run(global_step)) merged_summary_op = tf.summary.merge_all() real_images, synthetic_images = data_ops.loadData(dataset) # just use the last 500 synthetic for testing train_synthetic, test_synthetic = synthetic_images[: 10000], synthetic_images[ 10000:] real_len = len(real_images) synth_len = len(train_synthetic) num_test = len(test_synthetic) num_train = real_len + synth_len print 'train:', len(train_synthetic) print 'test:', len(test_synthetic) n_critic = 5
MAX_STEPS = a.MAX_STEPS CHECKPOINT_DIR = 'checkpoints/DATASET_' + DATASET + '/SCALE_' + str( SCALE) + '/NORM_' + str(NORM) + '/SELU_' + str(SELU) + '/' IMAGES_DIR = CHECKPOINT_DIR + 'images/' try: os.makedirs(IMAGES_DIR) except: pass # placeholders for data going into the network global_step = tf.Variable(0, name='global_step', trainable=False) z = tf.placeholder(tf.float32, shape=(BATCH_SIZE, 100), name='z') train_images_list = data_ops.loadData(DATA_DIR, DATASET) filename_queue = tf.train.string_input_producer(train_images_list) real_images = data_ops.read_input_queue(filename_queue, BATCH_SIZE) # generated images gen_images = netG(z, BATCH_SIZE) # get the output from D on the real and fake data errD_real = netD(real_images, BATCH_SIZE, SELU, NORM) errD_fake = netD(gen_images, BATCH_SIZE, SELU, NORM, reuse=True) # cost functions errD = tf.reduce_mean(errD_real) - tf.reduce_mean(errD_fake) errG = tf.reduce_mean(errD_fake) # gradient penalty
sys.path.insert(0, 'ops/') sys.path.insert(0, 'config/') import data_ops if __name__ == '__main__': CHECKPOINT_DIR = 'checkpoints/' IMAGES_DIR = CHECKPOINT_DIR+'images/' BATCH_SIZE=1 test_images = glob.glob(sys.argv[1]+'*.*') num_images = len(test_images) Data = data_ops.loadData(sys.argv[1], BATCH_SIZE, train=False) # The gray 'lightness' channel in range [-1, 1] gray_image = Data.inputs # The color channels in [-1, 1] range color_image = Data.targets # architecture from # http://hi.cs.waseda.ac.jp/~iizuka/projects/colorization/data/colorization_sig2016.pdf col_img = colorarch.architecture(gray_image, train=False) col_img = tf.image.convert_image_dtype(col_img, dtype=tf.uint8, saturate=True) saver = tf.train.Saver(max_to_keep=1) init = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer())
exp_pkl = open(CHECKPOINT_DIR + 'info.pkl', 'wb') data = pickle.dumps(exp_info) exp_pkl.write(data) exp_pkl.close() print print 'EPOCHS: ', EPOCHS print 'DATA_DIR: ', DATA_DIR print 'BATCH_SIZE: ', BATCH_SIZE print # global step that is saved with a model to keep track of how many steps/epochs global_step = tf.Variable(0, name='global_step', trainable=False) # load data Data = data_ops.loadData(DATA_DIR, BATCH_SIZE) num_train = Data.count gray_image = Data.inputs color_image = Data.targets # architecture from col_img = net.architecture(gray_image) #loss = tf.reduce_mean((ab_image-col_img)**2) loss = tf.reduce_mean(tf.nn.l2_loss(color_image - col_img)) train_op = tf.train.AdamOptimizer(learning_rate=1e-6).minimize( loss, global_step=global_step) saver = tf.train.Saver(max_to_keep=1) # tensorboard summaries