def test_generator_graph(self): for shape in ([4, 32, 32], [3, 128, 128], [2, 80, 400]): tf.reset_default_graph() img = tf.ones(shape + [3]) output_imgs = networks.generator(img) self.assertAllEqual(shape + [3], output_imgs.shape.as_list())
def main(): if IS_TRAINED: #initialize the model input_tf = tf.placeholder(tf.float32, [None, IMG_H, IMG_W, IMG_C]) train_phase = tf.placeholder(tf.bool) inpainting = generator("generator") patch_tf = inpainting(input_tf, train_phase) sess = tf.Session() sess.run(tf.global_variables_initializer()) saver = tf.train.Saver(tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, "generator")) saver.restore(sess, "./save_para/para.ckpt") #saver.restore(sess, "./save_para//.\\para.ckpt") #test the model file_path = "./data/ori_test/*.*" result_path ="./baseline_test" filenames = glob.glob(file_path) for i in range(len(filenames)): mask, X, Y = get_mask() img = misc.imresize(read_img_and_crop(filenames[i]), [IMG_H, IMG_W]) input = ((img * (1 - mask) + 255 * mask) / 127.5 - 1.0)[np.newaxis, :, :, :] patch = sess.run(patch_tf, feed_dict={input_tf: input, train_phase: False}) input[0, :, :, :][X:X + MASK_H, Y:Y + MASK_W, :] = patch[0, :, :, :] output_image = (input[0, :, :, :]+1)*127.5 input_image = (img * (1 - mask) + 255 * mask) #output = np.concatenate((img, mask*255, (input[0, :, :, :]+1)*127.5), 1) Image.fromarray(np.uint8(input_image)).save('/home/sunchenyu1993/CE7454/results/context_encoder_input/' + filenames[i].split('/')[-1]) Image.fromarray(np.uint8(output_image)).save('/home/sunchenyu1993/CE7454/results/context_encoder/' + filenames[i].split('/')[-1]) print ("saving: ",filenames[i].split('/')[-1] ) else: CE = ContextEncoder() CE.train()
def make_inference_graph(model_name, patch_dim): """Build the inference graph for either the X2Y or Y2X GAN. Args: model_name: The var scope name 'ModelX2Y' or 'ModelY2X'. patch_dim: An integer size of patches to feed to the generator. Returns: Tuple of (input_placeholder, generated_tensor). 知识点: tf.variable_scope: tf.get_variable(<name>, <shape>, <initializer>) 创建或返回给定名称的变量 tf.variable_scope(<scope_name>) 管理传给get_variable()的变量名称的作用域 tf.expand_dims: 想要维度增加一维,可以使用tf.expand_dims(input, dim, name=None)函数。 当然,我们常用tf.reshape(input, shape=[])也可以达到相同效果, 但是有些时候在构建图的过程中,placeholder没有被feed具体的值,这时就会包下面的错误:TypeError: Expected binary or unicode string, got 1 """ input_hwc_pl = tf.placeholder(tf.float32, [None, None, 3]) # Expand HWC to NHWC images_x = tf.expand_dims( data_provider.full_image_to_patch(input_hwc_pl, patch_dim), 0) with tf.variable_scope(model_name): with tf.variable_scope('Generator'): generated = networks.generator(images_x) return input_hwc_pl, generated
def get_GAN(shape,discriminator,generator,optimizer): discriminator.trainable = False inp = Input(shape=shape) x = generator(inp) out = discriminator(x) GAN = Model(inputs=inp,outputs=[x,out]) GAN.compile(loss=[vgg_loss,'binary_crossentropy'],loss_weights=[1,1e-3],optimizer=optimizer) return GAN
def train(): real_img = tf.placeholder(tf.float32, [None, H, W, 3]) label = tf.placeholder(tf.int32, [None]) z = tf.placeholder(tf.float32, [None, 100]) one_hot_label = tf.one_hot(label, NUMS_CLASS) labeled_z = tf.concat([z, one_hot_label], axis=1) G = generator("generator") D = discriminator("discriminator") fake_img = G(labeled_z) class_fake_logits, adv_fake_logits = D(fake_img, NUMS_CLASS) class_real_logits, adv_real_logits = D(real_img, NUMS_CLASS) loss_d_real = -tf.reduce_mean(tf.log(adv_real_logits + EPSILON)) loss_d_fake = -tf.reduce_mean(tf.log(1 - adv_fake_logits + EPSILON)) loss_cls_real = -tf.reduce_mean( tf.log( tf.reduce_sum(class_real_logits * one_hot_label, axis=1) + EPSILON)) loss_cls_fake = -tf.reduce_mean( tf.log( tf.reduce_sum(class_fake_logits * one_hot_label, axis=1) + EPSILON)) D_loss = loss_d_real + loss_d_fake + loss_cls_real G_loss = -tf.reduce_mean(tf.log(adv_fake_logits + EPSILON)) + loss_cls_fake D_opt = tf.train.AdamOptimizer(2e-4, beta1=0.5).minimize(D_loss, var_list=D.var_list()) G_opt = tf.train.AdamOptimizer(2e-4, beta1=0.5).minimize(G_loss, var_list=G.var_list()) sess = tf.Session() sess.run(tf.global_variables_initializer()) saver = tf.train.Saver() data, labels = read_face_data("./dataset/face_woman_man.mat") for i in range(50000): s = time.time() for j in range(1): BATCH, LABELS, Z = get_batch_face(data, labels, BATCHSIZE) BATCH = BATCH / 127.5 - 1.0 sess.run(D_opt, feed_dict={real_img: BATCH, label: LABELS, z: Z}) sess.run(G_opt, feed_dict={real_img: BATCH, label: LABELS, z: Z}) e = time.time() if i % 100 == 0: [D_LOSS, G_LOSS, FAKE_IMG] = sess.run([D_loss, G_loss, fake_img], feed_dict={ real_img: BATCH, label: LABELS, z: Z }) Image.fromarray(np.uint8((FAKE_IMG[0, :, :, :] + 1) * 127.5)).save("./results/" + str(i) + "_" + str(int(LABELS[0])) + ".jpg") print("Iteration: %d, D_loss: %f, G_loss: %f, update_time: %f" % (i, D_LOSS, G_LOSS, e - s)) if i % 500 == 0: saver.save(sess, "./save_para/model.ckpt") pass
def test_generator_grad_norm_progress(self): stable_stage_num_images = 2 transition_stage_num_images = 3 current_image_id_ph = tf.placeholder(tf.int32, []) progress = networks.compute_progress(current_image_id_ph, stable_stage_num_images, transition_stage_num_images, num_blocks=3) z = tf.random_normal([2, 10], dtype=tf.float32) x, _ = networks.generator( z, progress, _num_filters_stub, networks.ResolutionSchedule(start_resolutions=(4, 4), scale_base=2, num_resolutions=3)) fake_loss = tf.reduce_sum(tf.square(x)) grad_norms = [ _get_grad_norm( fake_loss, tf.trainable_variables('.*/progressive_gan_block_1/.*')), _get_grad_norm( fake_loss, tf.trainable_variables('.*/progressive_gan_block_2/.*')), _get_grad_norm( fake_loss, tf.trainable_variables('.*/progressive_gan_block_3/.*')) ] grad_norms_output = None with self.test_session(use_gpu=True) as sess: sess.run(tf.global_variables_initializer()) x1_np = sess.run(x, feed_dict={current_image_id_ph: 0.12}) x2_np = sess.run(x, feed_dict={current_image_id_ph: 1.8}) grad_norms_output = np.array([ sess.run(grad_norms, feed_dict={current_image_id_ph: i}) for i in range(15) # total num of images ]) self.assertEqual((2, 16, 16, 3), x1_np.shape) self.assertEqual((2, 16, 16, 3), x2_np.shape) # The gradient of block_1 is always on. self.assertEqual( np.argmax(grad_norms_output[:, 0] > 0), 0, 'gradient norms {} for block 1 is not always on'.format( grad_norms_output[:, 0])) # The gradient of block_2 is on after 1 stable stage. self.assertEqual( np.argmax(grad_norms_output[:, 1] > 0), 3, 'gradient norms {} for block 2 is not on at step 3'.format( grad_norms_output[:, 1])) # The gradient of block_3 is on after 2 stable stage + 1 transition stage. self.assertEqual( np.argmax(grad_norms_output[:, 2] > 0), 8, 'gradient norms {} for block 3 is not on at step 8'.format( grad_norms_output[:, 2]))
def train(epochs=1,batch_size=128): if(not os.path.isdir('/content/output/')): os.mkdir('/content/output/') downscale_factor = 4 batch_count = int(x_train_hr.shape[0] / batch_size) shape = (img_shape[0]//downscale_factor, img_shape[1]//downscale_factor, img_shape[2]) gen = generator(shape).generator_model() disc = discriminator(img_shape).discriminator_model() adam = Adam(lr=1E-4, beta_1=0.9, beta_2=0.999, epsilon=1e-08) gen.compile(loss=vgg_loss, optimizer=adam) disc.compile(loss="binary_crossentropy", optimizer=adam) shape = (img_shape[0]//downscale_factor, img_shape[1]//downscale_factor, 3) gan = get_GAN(shape, disc, gen, adam) for e in range(1, epochs+1): print ('-'*15, 'Epoch %d' % e, '-'*15) for _ in range(batch_count): rand_nums = np.random.randint(0, x_train_hr.shape[0], size=batch_size) image_batch_hr = x_train_hr[rand_nums] image_batch_lr = x_train_lr[rand_nums] generated_images_sr = gen.predict(image_batch_lr) real_data_Y = np.ones(batch_size) - np.random.random_sample(batch_size)*0.2 fake_data_Y = np.random.random_sample(batch_size)*0.2 disc.trainable = True d_loss_real = disc.train_on_batch(image_batch_hr, real_data_Y) d_loss_fake = disc.train_on_batch(generated_images_sr, fake_data_Y) #d_loss = 0.5 * np.add(d_loss_fake, d_loss_real) rand_nums = np.random.randint(0, x_train_hr.shape[0], size=batch_size) image_batch_hr = x_train_hr[rand_nums] image_batch_lr = x_train_lr[rand_nums] gan_Y = np.ones(batch_size) - np.random.random_sample(batch_size)*0.2 disc.trainable = False loss_gan = gan.train_on_batch(image_batch_lr, [image_batch_hr,gan_Y]) print("Loss HR , Loss LR, Loss GAN") print(d_loss_real, d_loss_fake, loss_gan) if e == 1 or e % 5 == 0: plot_generated_images(e, gen) if e % 50 == 0: generator.save('./output/gen_model%d.h5' % e) discriminator.save('./output/dis_model%d.h5' % e) gan.save('./output/gan_model%d.h5' % e) train(20000,4)
def test_generator(self): tf.set_random_seed(1234) batch_size = 100 noise = tf.random_normal([batch_size, 64]) image = networks.generator(noise) with self.test_session(use_gpu=True) as sess: sess.run(tf.global_variables_initializer()) image_np = image.eval() self.assertAllEqual([batch_size, 32, 32, 3], image_np.shape) self.assertTrue(np.all(np.abs(image_np) <= 1))
def _generator_fn(z): """Builds generator network.""" return networks.generator( z, progress, _num_filters_fn, resolution_schedule, num_blocks=num_blocks, kernel_size=kernel_size, colors=colors, to_rgb_activation=(tf.tanh if kwargs['to_rgb_use_tanh_activation'] else None))
def make_inference_graph(model_name, patch_dim): """Build the inference graph for either the X2Y or Y2X GAN. Args: model_name: The var scope name 'ModelX2Y' or 'ModelY2X'. patch_dim: An integer size of patches to feed to the generator. Returns: Tuple of (input_placeholder, generated_tensor). """ input_hwc_pl = tf.placeholder(tf.float32, [None, None, 3]) # Expand HWC to NHWC images_x = tf.expand_dims( data_provider.full_image_to_patch(input_hwc_pl, patch_dim), 0) with tf.variable_scope(model_name): with tf.variable_scope('Generator'): generated = networks.generator(images_x) return input_hwc_pl, generated
def generate_fixed_z(): label = tf.placeholder(tf.float32, [None, NUMS_CLASS]) z = tf.placeholder(tf.float32, [None, 100]) labeled_z = tf.concat([z, label], axis=1) G = generator("generator") fake_img = G(labeled_z) sess = tf.Session() sess.run(tf.global_variables_initializer()) saver = tf.train.Saver() saver.restore(sess, "./save_para/model.ckpt") LABELS, Z = label_from_0_to_1() if not os.path.exists("./generate_fixed_noise"): os.mkdir("./generate_fixed_noise") FAKE_IMG = sess.run(fake_img, feed_dict={label: LABELS, z: Z}) for i in range(10): Image.fromarray(np.uint8((FAKE_IMG[i, :, :, :] + 1) * 127.5)).save("./generate_fixed_noise/" + str(i) + ".jpg")
def generate_fixed_label(): label = tf.placeholder(tf.int32, [None]) z = tf.placeholder(tf.float32, [None, 100]) one_hot_label = tf.one_hot(label, NUMS_CLASS) labeled_z = tf.concat([z, one_hot_label], axis=1) G = generator("generator") fake_img = G(labeled_z) sess = tf.Session() sess.run(tf.global_variables_initializer()) saver = tf.train.Saver() saver.restore(sess, "./save_para/model.ckpt") Z = from_noise0_to_noise1() LABELS = np.ones([10])#woman: LABELS = np.ones([10]), man: LABELS = np.zeros([10]) if not os.path.exists("./generate_fixed_label"): os.mkdir("./generate_fixed_label") FAKE_IMG = sess.run(fake_img, feed_dict={label: LABELS, z: Z}) for i in range(10): Image.fromarray(np.uint8((FAKE_IMG[i, :, :, :] + 1) * 127.5)).save("./generate_fixed_label/" + str(i) + "_" + str(int(LABELS[i])) + ".jpg")
def __init__(self, opts): super(DivCo_DCGAN, self).__init__() # parameters lr = 0.0002 self.nz = opts.nz self.opt = opts self.class_num = opts.class_num self.G = networks.generator(opts) self.D = networks.discriminator(opts) self.gen_opt = torch.optim.Adam(self.G.parameters(), lr=lr, betas=(0.5, 0.999), weight_decay=0.0001) self.dis_opt = torch.optim.Adam(self.D.parameters(), lr=lr, betas=(0.5, 0.999), weight_decay=0.0001) self.BCE_loss = torch.nn.BCELoss() self.cross_entropy_loss = torch.nn.CrossEntropyLoss()
def main(checkpoint_dir, grid_size, latent_dim, output_name): """Generates images from a pretrained DCGAN.""" if grid_size > 20: raise ValueError( 'To prevent the saved image from getting too large, the grid size cannot be more than 20. ' 'To save more than 400 images, run this script multiple times with different output names.' ) num_images = grid_size**2 if output_name is None: output_name = 'generated_images_{}.png'.format(num_images) assert output_name.endswith( '.png'), 'Saved file must have a PNG extension.' # Get checkpoint path checkpoint_path = tf.train.latest_checkpoint(checkpoint_dir) if checkpoint_path is None: print('No checkpoint could be found. Exiting.') exit(0) # Generate new images with tf.variable_scope('Generator'): images = networks.generator(tf.random_normal([num_images, latent_dim]), is_training=False) # Tile them in a single image reshaped_images = tf.contrib.gan.eval.image_reshaper(images, num_cols=grid_size) uint8_images = float_image_to_uint8(reshaped_images) image_write_ops = tf.write_file( os.path.join(checkpoint_dir, 'eval', output_name), tf.image.encode_png(uint8_images[0])) # Restore checkpoint and run ops session_creator = tf.train.ChiefSessionCreator( checkpoint_filename_with_path=checkpoint_path) with tf.train.MonitoredSession(session_creator=session_creator) as sess: sess.run(image_write_ops)
def __init__(self): # Paper: Context Encoders: Feature Learning by Inpainting self.inputs = tf.placeholder(tf.float32, [None, IMG_H, IMG_W, IMG_C]) self.patch = tf.placeholder(tf.float32, [None, MASK_H, MASK_W, IMG_C]) self.train_phase = tf.placeholder(tf.bool) G = generator("generator") D = discriminator("discriminator") self.patch_fake = G(self.inputs, self.train_phase) self.fake_logits = D(self.patch_fake, self.train_phase) self.real_logits = D(self.patch, self.train_phase) self.D_loss = -tf.reduce_mean( tf.log(self.real_logits + EPSILON) + tf.log(1 - self.fake_logits + EPSILON)) self.G_loss = -tf.reduce_mean( tf.log(self.fake_logits + EPSILON)) + 100 * tf.reduce_mean( tf.reduce_sum(tf.square(self.patch - self.patch_fake), [1, 2, 3])) self.D_Opt = tf.train.AdamOptimizer(2e-4).minimize( self.D_loss, var_list=D.get_var()) self.G_Opt = tf.train.AdamOptimizer(2e-4).minimize( self.G_loss, var_list=G.get_var()) self.sess = tf.Session() self.sess.run(tf.global_variables_initializer())
def main(): if IS_TRAINED: #initialize the model input_tf = tf.placeholder(tf.float32, [None, IMG_H, IMG_W, IMG_C]) train_phase = tf.placeholder(tf.bool) inpainting = generator("generator") patch_tf = inpainting(input_tf, train_phase) sess = tf.Session() sess.run(tf.global_variables_initializer()) saver = tf.train.Saver(tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, "generator")) saver.restore(sess, "./save_para//.\\para.ckpt") #test the model file_path = "./cats_bigger_than_128x128/00000003_020.jpg" mask, X, Y = get_mask() img = misc.imresize(read_img_and_crop(file_path), [IMG_H, IMG_W]) input = ((img * (1 - mask) + 255 * mask) / 127.5 - 1.0)[np.newaxis, :, :, :] patch = sess.run(patch_tf, feed_dict={input_tf: input, train_phase: False}) input[0, :, :, :][X:X + MASK_H, Y:Y + MASK_W, :] = patch[0, :, :, :] output = np.concatenate((img, mask*255, (input[0, :, :, :]+1)*127.5), 1) Image.fromarray(np.uint8(output)).show() else: CE = ContextEncoder() CE.train()
def main(model_dir): # The placehodler for feeding input noise to the generator. z_placeholder = tf.placeholder(tf.float32, [None, Z_DIMENSIONS], name='z_placeholder') # The generator. g_z = generator(z_placeholder, BATCH_SIZE, Z_DIMENSIONS) with tf.Session() as sess: ckpt = tf.train.latest_checkpoint(model_dir) saver = tf.train.Saver() saver.restore(sess, ckpt) z_batch = np.random.normal(0, 1, size=[1, Z_DIMENSIONS]) generated_img = sess.run(g_z, { z_placeholder: z_batch })[0] generated_img = generated_img * 128 generated_img = generated_img + 127 generated_img = generated_img.astype(np.uint8) cv2.imshow('Generated Image', generated_img) cv2.waitKey(0)
def generator_fn_specgram(inputs, **kwargs): """Builds generator network.""" # inputs = (noises, one_hot_labels) with tf.variable_scope('generator_cond'): z = tf.concat(inputs, axis=1) if kwargs['to_rgb_activation'] == 'tanh': to_rgb_activation = tf.tanh elif kwargs['to_rgb_activation'] == 'linear': to_rgb_activation = lambda x: x fake_images, end_points = networks.generator( z, kwargs['progress'], lambda block_id: _num_filters_fn(block_id, **kwargs), kwargs['resolution_schedule'], num_blocks=kwargs['num_blocks'], kernel_size=kwargs['kernel_size'], colors=2, to_rgb_activation=to_rgb_activation, simple_arch=kwargs['simple_arch']) shape = fake_images.shape normalizer = data_normalizer.registry[kwargs['data_normalizer']](kwargs) fake_images = normalizer.denormalize_op(fake_images) fake_images.set_shape(shape) return fake_images, end_points
def main(gan): # Load anime data set. anime = read_data_set() # The placehodler for feeding input noise to the generator. z_placeholder = tf.placeholder(tf.float32, [None, Z_DIMENSIONS], name='z_placeholder') # The placehodler for feeding input images to the discriminator. x_placeholder = tf.placeholder(tf.float32, shape=[None, 64, 64, 3], name='x_placeholder') if gan == GAN.BASIC: ranged_output = True elif gan == GAN.LSGAN: ranged_output = True # The generated images. g_z = generator(z_placeholder, BATCH_SIZE, Z_DIMENSIONS) # The discriminator prediction probability for the real images. d_x = discriminator(x_placeholder, ranged_output) # The discriminator prediction probability for the generated images. d_g = discriminator(g_z, ranged_output, reuse_variables=True) if gan == GAN.BASIC: loss = BasicLoss(d_x, d_g) elif gan == GAN.LSGAN: loss = LeastSquaresLoss(d_x, d_g, 0, 1, 1) # Two Loss Functions for discriminator. d_loss_real = loss.d_loss_real() d_loss_fake = loss.d_loss_fake() # Loss function for generator. g_loss = loss.g_loss() # Get the varaibles for different network. tvars = tf.trainable_variables() d_vars = [var for var in tvars if 'd_' in var.name] g_vars = [var for var in tvars if 'g_' in var.name] # Train the discriminator. d_trainer_fake = tf.train.AdamOptimizer(0.0003).minimize(d_loss_fake, var_list=d_vars) d_trainer_real = tf.train.AdamOptimizer(0.0003).minimize(d_loss_real, var_list=d_vars) # Train the generator. g_trainer = tf.train.AdamOptimizer(0.0001).minimize(g_loss, var_list=g_vars) # From this point forward, reuse variables. tf.get_variable_scope().reuse_variables() with tf.Session() as sess: # Send summary statistics to TensorBoard. tf.summary.scalar('Generator_loss', g_loss) tf.summary.scalar('Discriminator_loss_real', d_loss_real) tf.summary.scalar('Discriminator_loss_fake', d_loss_fake) folder_id = datetime.datetime.now().strftime('%Y%m%d-%H%M%S') folder_root = os.path.join(RESULT_ROOT, folder_id) images_for_tensorboard = generator(z_placeholder, BATCH_SIZE, Z_DIMENSIONS) tf.summary.image('Generated_images', images_for_tensorboard, 5) merged = tf.summary.merge_all() logdir = os.path.join(folder_root, 'tensorboard') writer = tf.summary.FileWriter(logdir, sess.graph) model_dir = os.path.join(folder_root, 'model') model_path = '{}/model.ckpt'.format(model_dir) if not os.path.exists(model_dir): os.makedirs(model_dir) img_dir = os.path.join(folder_root, 'images') if not os.path.exists(img_dir): os.makedirs(img_dir) saver = tf.train.Saver() sess.run(tf.global_variables_initializer()) # Pre-train discriminator. for i in range(300): z_batch = np.random.normal(0, 1, size=[BATCH_SIZE, Z_DIMENSIONS]) real_image_batch = anime.sample(BATCH_SIZE) _, __ = sess.run( [d_trainer_real, d_trainer_fake], {x_placeholder: real_image_batch, z_placeholder: z_batch}, ) # Train generator and discriminator together for i in range(ITERATION): real_image_batch = anime.sample(BATCH_SIZE) z_batch = np.random.normal(0, 1, size=[BATCH_SIZE, Z_DIMENSIONS]) # Train discriminator on both real and fake images. for j in range(D_UPDATE): _, __, d_loss_real_score, d_loss_fake_score = sess.run( [d_trainer_real, d_trainer_fake, d_loss_real, d_loss_fake], {x_placeholder: real_image_batch, z_placeholder: z_batch}, ) # Train generator. for j in range(G_UPDATE): z_batch = np.random.normal(0, 1, size=[BATCH_SIZE, Z_DIMENSIONS]) _, g_loss_score = sess.run([g_trainer, g_loss], feed_dict={z_placeholder: z_batch}) if i % 10 == 0: # Update TensorBoard with summary statistics. z_batch = np.random.normal(0, 1, size=[BATCH_SIZE, Z_DIMENSIONS]) summary = sess.run( merged, {z_placeholder: z_batch, x_placeholder: real_image_batch}, ) writer.add_summary(summary, i) if i % 100 == 0: z_batch = np.random.normal(0, 1, size=[1, Z_DIMENSIONS]) generated_img = sess.run(g_z, { z_placeholder: z_batch }) generated_img = generated_img.reshape([64, 64, 3]) generated_img = generated_img * 128 generated_img = generated_img + 127 generated_img = generated_img.astype(np.uint8) cv2.imwrite('{}/{:05d}.png'.format(img_dir, i), generated_img) if i % 100 == 0: print("d_loss_real_score:", d_loss_real_score, "d_loss_fake_score:", d_loss_fake_score, "g_loss_score:", g_loss_score) saver.save(sess, model_path)
help='pre_trained cartoongan model path') parser.add_argument('--image_dir', required=True, default='image_dir', help='test image path') parser.add_argument('--output_image_dir', required=True, default='output_image_dir', help='output test image path') args = parser.parse_args() device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') if torch.backends.cudnn.enabled: torch.backends.cudnn.benchmark = True G = networks.generator(args.in_ngc, args.out_ngc, args.ngf, args.nb) if torch.cuda.is_available(): G.load_state_dict(torch.load(args.pre_trained_model)) else: # cpu mode G.load_state_dict( torch.load(args.pre_trained_model, map_location=lambda storage, loc: storage)) G.to(device) src_transform = transforms.Compose([ transforms.Resize((args.input_size_h, args.input_size_w)), transforms.ToTensor(), transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)) ]) # utils.data_load(os.path.join('data', args.src_data), 'test', src_transform, 1, shuffle=True, drop_last=True)
parser.add_argument('--pre_train_epoch', type=int, default=10) parser.add_argument('--lrD', type=float, default=0.0002, help='learning rate, default=0.0002') parser.add_argument('--lrG', type=float, default=0.0002, help='learning rate, default=0.0002') parser.add_argument('--con_lambda', type=float, default=10, help='lambda for content loss') parser.add_argument('--beta1', type=float, default=0.5, help='beta1 for Adam optimizer') parser.add_argument('--beta2', type=float, default=0.999, help='beta2 for Adam optimizer') parser.add_argument('--pre_trained_model', required=True, default='pre_trained_model', help='pre_trained cartoongan model path') parser.add_argument('--image_dir', required=True, default='image_dir', help='test image path') parser.add_argument('--output_image_dir', required=True, default='output_image_dir', help='output test image path') args = parser.parse_args() device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') if torch.backends.cudnn.enabled: torch.backends.cudnn.benchmark = True G = networks.generator(args.in_ngc, args.out_ngc, args.ngf, args.nb) if torch.cuda.is_available(): G.load_state_dict(torch.load(args.pre_trained_model)) else: # cpu mode G.load_state_dict(torch.load(args.pre_trained_model, map_location=lambda storage, loc: storage)) G.to(device) src_transform = transforms.Compose([ transforms.Resize((args.input_size, args.input_size)), transforms.ToTensor(), transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)) ]) # utils.data_load(os.path.join('data', args.src_data), 'test', src_transform, 1, shuffle=True, drop_last=True) image_src = utils.data_load(os.path.join(args.image_dir), 'test', src_transform, 1, shuffle=True, drop_last=True)
def test_generator_run_multi_channel(self): img_batch = tf.zeros([3, 128, 128, 5]) model_output = networks.generator(img_batch) with self.test_session() as sess: sess.run(tf.global_variables_initializer()) sess.run(model_output)
def test_generator_invalid_channels(self): with self.assertRaisesRegexp( ValueError, 'Last dimension shape must be known but is None'): img = tf.placeholder(tf.float32, shape=[4, 32, 32, None]) networks.generator(img)
def test_generator_graph_unknown_batch_dim(self): img = tf.placeholder(tf.float32, shape=[None, 32, 32, 3]) output_imgs = networks.generator(img) self.assertAllEqual([None, 32, 32, 3], output_imgs.shape.as_list())
target_size = (64, 64, 1) x_m = [] for i in range(len(X_train)): x = cv2.resize(X_train[i], target_size[:2]) x = x.reshape(target_size[0], target_size[1], target_size[2]) x_m.append(x) X_train = np.asarray(x_m) noise_dim = 100 n_critic = 5 n_labels = 10 n_channels = 3 gen = generator(target_size[0], target_size[1], 512, noise_dim, n_labels, target_size[2]) disc = discriminator(target_size[0], target_size[1], 256, n_labels, target_size[2], wgan=True) opt = Adam(0.0002, 0.5) #------------------------------- # Construct Computational Graph # for the Critic #------------------------------- # Freeze generator's layers while training critic gen.trainable = False # Image input (real sample)
def test_generator_run(self): img_batch = tf.zeros([3, 128, 128, 3]) model_output = networks.generator(img_batch) with self.test_session() as sess: sess.run(tf.global_variables_initializer()) sess.run(model_output)
help='pre_trained cartoongan model path') parser.add_argument('--image_dir', required=True, default='image_dir', help='test image path') parser.add_argument('--output_image_dir', required=True, default='output_image_dir', help='output test image path') args = parser.parse_args() device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') if torch.backends.cudnn.enabled: torch.backends.cudnn.benchmark = True G = networks.generator() if torch.cuda.is_available(): G.load_state_dict(torch.load(args.pre_trained_model)) else: # cpu mode G.load_state_dict( torch.load(args.pre_trained_model, map_location=lambda storage, loc: storage)) G.to(device) src_transform = transforms.Compose([ transforms.Resize((args.input_size, args.input_size)), transforms.ToTensor(), transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)) ]) # utils.data_load(os.path.join('data', args.src_data), 'test', src_transform, 1, shuffle=True, drop_last=True)
def main(): device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') if torch.backends.cudnn.enabled: torch.backends.cudnn.benchmark = True prepare_result() make_edge_promoting_img() # data_loader src_transform = transforms.Compose([ transforms.Resize((args.input_size, args.input_size)), transforms.ToTensor(), transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)) ]) tgt_transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)) ]) train_loader_src = utils.data_load(os.path.join('data', args.src_data), 'train', src_transform, args.batch_size, shuffle=True, drop_last=True) train_loader_tgt = utils.data_load(os.path.join('data', args.tgt_data), 'pair', tgt_transform, args.batch_size, shuffle=True, drop_last=True) test_loader_src = utils.data_load(os.path.join('data', args.src_data), 'test', src_transform, 1, shuffle=True, drop_last=True) # network G = networks.generator(args.in_ngc, args.out_ngc, args.ngf, args.nb) if args.latest_generator_model != '': if torch.cuda.is_available(): G.load_state_dict(torch.load(args.latest_generator_model)) else: # cpu mode G.load_state_dict( torch.load(args.latest_generator_model, map_location=lambda storage, loc: storage)) D = networks.discriminator(args.in_ndc, args.out_ndc, args.ndf) if args.latest_discriminator_model != '': if torch.cuda.is_available(): D.load_state_dict(torch.load(args.latest_discriminator_model)) else: D.load_state_dict( torch.load(args.latest_discriminator_model, map_location=lambda storage, loc: storage)) VGG = networks.VGG19(init_weights=args.vgg_model, feature_mode=True) G.to(device) D.to(device) VGG.to(device) G.train() D.train() VGG.eval() print('---------- Networks initialized -------------') utils.print_network(G) utils.print_network(D) utils.print_network(VGG) print('-----------------------------------------------') # loss BCE_loss = nn.BCELoss().to(device) L1_loss = nn.L1Loss().to(device) # Adam optimizer G_optimizer = optim.Adam(G.parameters(), lr=args.lrG, betas=(args.beta1, args.beta2)) D_optimizer = optim.Adam(D.parameters(), lr=args.lrD, betas=(args.beta1, args.beta2)) G_scheduler = optim.lr_scheduler.MultiStepLR( optimizer=G_optimizer, milestones=[args.train_epoch // 2, args.train_epoch // 4 * 3], gamma=0.1) D_scheduler = optim.lr_scheduler.MultiStepLR( optimizer=D_optimizer, milestones=[args.train_epoch // 2, args.train_epoch // 4 * 3], gamma=0.1) pre_train_hist = {} pre_train_hist['Recon_loss'] = [] pre_train_hist['per_epoch_time'] = [] pre_train_hist['total_time'] = [] """ Pre-train reconstruction """ if args.latest_generator_model == '': print('Pre-training start!') start_time = time.time() for epoch in range(args.pre_train_epoch): epoch_start_time = time.time() Recon_losses = [] for x, _ in train_loader_src: x = x.to(device) # train generator G G_optimizer.zero_grad() x_feature = VGG((x + 1) / 2) G_ = G(x) G_feature = VGG((G_ + 1) / 2) Recon_loss = 10 * L1_loss(G_feature, x_feature.detach()) Recon_losses.append(Recon_loss.item()) pre_train_hist['Recon_loss'].append(Recon_loss.item()) Recon_loss.backward() G_optimizer.step() per_epoch_time = time.time() - epoch_start_time pre_train_hist['per_epoch_time'].append(per_epoch_time) print('[%d/%d] - time: %.2f, Recon loss: %.3f' % ((epoch + 1), args.pre_train_epoch, per_epoch_time, torch.mean(torch.FloatTensor(Recon_losses)))) total_time = time.time() - start_time pre_train_hist['total_time'].append(total_time) with open(os.path.join(args.name + '_results', 'pre_train_hist.pkl'), 'wb') as f: pickle.dump(pre_train_hist, f) with torch.no_grad(): G.eval() for n, (x, _) in enumerate(train_loader_src): x = x.to(device) G_recon = G(x) result = torch.cat((x[0], G_recon[0]), 2) path = os.path.join( args.name + '_results', 'Reconstruction', args.name + '_train_recon_' + str(n + 1) + '.png') plt.imsave(path, (result.cpu().numpy().transpose(1, 2, 0) + 1) / 2) if n == 4: break for n, (x, _) in enumerate(test_loader_src): x = x.to(device) G_recon = G(x) result = torch.cat((x[0], G_recon[0]), 2) path = os.path.join( args.name + '_results', 'Reconstruction', args.name + '_test_recon_' + str(n + 1) + '.png') plt.imsave(path, (result.cpu().numpy().transpose(1, 2, 0) + 1) / 2) if n == 4: break else: print('Load the latest generator model, no need to pre-train') train_hist = {} train_hist['Disc_loss'] = [] train_hist['Gen_loss'] = [] train_hist['Con_loss'] = [] train_hist['per_epoch_time'] = [] train_hist['total_time'] = [] print('training start!') start_time = time.time() real = torch.ones(args.batch_size, 1, args.input_size // 4, args.input_size // 4).to(device) fake = torch.zeros(args.batch_size, 1, args.input_size // 4, args.input_size // 4).to(device) for epoch in range(args.train_epoch): epoch_start_time = time.time() G.train() Disc_losses = [] Gen_losses = [] Con_losses = [] for (x, _), (y, _) in zip(train_loader_src, train_loader_tgt): e = y[:, :, :, args.input_size:] y = y[:, :, :, :args.input_size] x, y, e = x.to(device), y.to(device), e.to(device) # train D D_optimizer.zero_grad() D_real = D(y) D_real_loss = BCE_loss(D_real, real) G_ = G(x) D_fake = D(G_) D_fake_loss = BCE_loss(D_fake, fake) D_edge = D(e) D_edge_loss = BCE_loss(D_edge, fake) Disc_loss = D_real_loss + D_fake_loss + D_edge_loss Disc_losses.append(Disc_loss.item()) train_hist['Disc_loss'].append(Disc_loss.item()) Disc_loss.backward() D_optimizer.step() # train G G_optimizer.zero_grad() G_ = G(x) D_fake = D(G_) D_fake_loss = BCE_loss(D_fake, real) x_feature = VGG((x + 1) / 2) G_feature = VGG((G_ + 1) / 2) Con_loss = args.con_lambda * L1_loss(G_feature, x_feature.detach()) Gen_loss = D_fake_loss + Con_loss Gen_losses.append(D_fake_loss.item()) train_hist['Gen_loss'].append(D_fake_loss.item()) Con_losses.append(Con_loss.item()) train_hist['Con_loss'].append(Con_loss.item()) Gen_loss.backward() G_optimizer.step() G_scheduler.step() D_scheduler.step() per_epoch_time = time.time() - epoch_start_time train_hist['per_epoch_time'].append(per_epoch_time) print( '[%d/%d] - time: %.2f, Disc loss: %.3f, Gen loss: %.3f, Con loss: %.3f' % ((epoch + 1), args.train_epoch, per_epoch_time, torch.mean(torch.FloatTensor(Disc_losses)), torch.mean(torch.FloatTensor(Gen_losses)), torch.mean(torch.FloatTensor(Con_losses)))) if epoch % 2 == 1 or epoch == args.train_epoch - 1: with torch.no_grad(): G.eval() for n, (x, _) in enumerate(train_loader_src): x = x.to(device) G_recon = G(x) result = torch.cat((x[0], G_recon[0]), 2) path = os.path.join( args.name + '_results', 'Transfer', str(epoch + 1) + '_epoch_' + args.name + '_train_' + str(n + 1) + '.png') plt.imsave(path, (result.cpu().numpy().transpose(1, 2, 0) + 1) / 2) if n == 4: break for n, (x, _) in enumerate(test_loader_src): x = x.to(device) G_recon = G(x) result = torch.cat((x[0], G_recon[0]), 2) path = os.path.join( args.name + '_results', 'Transfer', str(epoch + 1) + '_epoch_' + args.name + '_test_' + str(n + 1) + '.png') plt.imsave(path, (result.cpu().numpy().transpose(1, 2, 0) + 1) / 2) if n == 4: break torch.save( G.state_dict(), os.path.join(args.name + '_results', 'generator_latest.pkl')) torch.save( D.state_dict(), os.path.join(args.name + '_results', 'discriminator_latest.pkl')) total_time = time.time() - start_time train_hist['total_time'].append(total_time) print("Avg one epoch time: %.2f, total %d epochs time: %.2f" % (torch.mean(torch.FloatTensor( train_hist['per_epoch_time'])), args.train_epoch, total_time)) print("Training finish!... save training results") torch.save(G.state_dict(), os.path.join(args.name + '_results', 'generator_param.pkl')) torch.save(D.state_dict(), os.path.join(args.name + '_results', 'discriminator_param.pkl')) with open(os.path.join(args.name + '_results', 'train_hist.pkl'), 'wb') as f: pickle.dump(train_hist, f)
def test_generator_invalid_input(self): with self.assertRaisesRegexp(ValueError, 'must have rank 4'): networks.generator(tf.zeros([28, 28, 3]))
def main(): # Load mnist data. mnist = input_data.read_data_sets('MNIST_data/') # The placehodler for feeding input noise to the generator. z_placeholder = tf.placeholder(tf.float32, [None, Z_DIMENSIONS], name='z_placeholder') # The placehodler for feeding input images to the discriminator. x_placeholder = tf.placeholder(tf.float32, shape=[None, 28, 28, 1], name='x_placeholder') # The generated images. Gz = generator(z_placeholder, BATCH_SIZE, Z_DIMENSIONS) # The discriminator prediction probability for the real images. Dx = discriminator(x_placeholder) # The discriminator prediction probability for the generated images. Dg = discriminator(Gz, reuse_variables=True) # Two Loss Functions for discriminator. d_loss_real = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits(logits=Dx, labels=tf.ones_like(Dx))) d_loss_fake = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits(logits=Dg, labels=tf.zeros_like(Dg))) # Loss function for generator. g_loss = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits(logits=Dg, labels=tf.ones_like(Dg))) # Get the varaibles for different network. tvars = tf.trainable_variables() d_vars = [var for var in tvars if 'd_' in var.name] g_vars = [var for var in tvars if 'g_' in var.name] # Train the discriminator. d_trainer_fake = tf.train.AdamOptimizer(0.0003).minimize(d_loss_fake, var_list=d_vars) d_trainer_real = tf.train.AdamOptimizer(0.0003).minimize(d_loss_real, var_list=d_vars) # Train the generator. g_trainer = tf.train.AdamOptimizer(0.0001).minimize(g_loss, var_list=g_vars) # From this point forward, reuse variables. tf.get_variable_scope().reuse_variables() with tf.Session() as sess: # Send summary statistics to TensorBoard. tf.summary.scalar('Generator_loss', g_loss) tf.summary.scalar('Discriminator_loss_real', d_loss_real) tf.summary.scalar('Discriminator_loss_fake', d_loss_fake) images_for_tensorboard = generator(z_placeholder, BATCH_SIZE, Z_DIMENSIONS) tf.summary.image('Generated_images', images_for_tensorboard, 5) merged = tf.summary.merge_all() logdir = 'tensorboard/{}/'.format( datetime.datetime.now().strftime('%Y%m%d-%H%M%S')) writer = tf.summary.FileWriter(logdir, sess.graph) sess.run(tf.global_variables_initializer()) # Pre-train discriminator. for i in range(300): z_batch = np.random.normal(0, 1, size=[BATCH_SIZE, Z_DIMENSIONS]) real_image_batch = mnist.train.next_batch(BATCH_SIZE)[0].reshape( [BATCH_SIZE, 28, 28, 1]) _, __, dLossReal, dLossFake = sess.run( [d_trainer_real, d_trainer_fake, d_loss_real, d_loss_fake], { x_placeholder: real_image_batch, z_placeholder: z_batch }) if i % 100 == 0: print("dLossReal:", dLossReal, "dLossFake:", dLossFake) # Train generator and discriminator together for i in range(100000): real_image_batch = mnist.train.next_batch(BATCH_SIZE)[0].reshape( [BATCH_SIZE, 28, 28, 1]) z_batch = np.random.normal(0, 1, size=[BATCH_SIZE, Z_DIMENSIONS]) # Train discriminator on both real and fake images. _, __, dLossReal, dLossFake = sess.run( [d_trainer_real, d_trainer_fake, d_loss_real, d_loss_fake], { x_placeholder: real_image_batch, z_placeholder: z_batch }) # Train generator. z_batch = np.random.normal(0, 1, size=[BATCH_SIZE, Z_DIMENSIONS]) _ = sess.run(g_trainer, feed_dict={z_placeholder: z_batch}) if i % 10 == 0: # Update TensorBoard with summary statistics. z_batch = np.random.normal(0, 1, size=[BATCH_SIZE, Z_DIMENSIONS]) summary = sess.run(merged, { z_placeholder: z_batch, x_placeholder: real_image_batch }) writer.add_summary(summary, i)
#writer7 = threading.Thread(target=producer, args=(q, stop_event, train_paths, y_train, batch_size, target_size, prep_func)) #writer8 = threading.Thread(target=producer, args=(q, stop_event, train_paths, y_train, batch_size, target_size, prep_func)) writer.start() writer2.start() writer3.start() writer4.start() writer5.start() writer6.start() #writer7.start() #writer8.start() # gen = generator(target_size[0], target_size[1], 1024, noise_dim, n_labels, target_size[2], tanh=True) disc = discriminator(target_size[0], target_size[1], 512, n_labels, target_size[2], wgan=True) opt = Adam(0.0002, 0.5) disc.compile(loss=['binary_crossentropy', null_loss()], optimizer=opt, metrics=['accuracy']) frozen_disc = Model(inputs=disc.inputs, outputs=disc.outputs) frozen_disc.trainable = False
d_trainer = tf.train.AdamOptimizer(LEFTOUT_RATE).minimize(d_loss, var_list=d_vars) g_trainer = tf.train.AdamOptimizer(LEFTOUT_RATE).minimize(g_loss, var_list=g_vars) ### load training dataset training_set = mnist.read_data_sets("MNIST_data") samples = [] with tf.Session() as s: s.run(tf.global_variables_initializer()) for epoch in range(EPOCHS): batches = training_set.train.num_examples // BATCH_SIZE for i in range(batches): batch = training_set.train.next_batch(BATCH_SIZE) batch_images = batch[0].reshape(BATCH_SIZE, 784) batch_images = batch_images * 2 - 1 batch_seed = np.random.uniform(-1, 1, size=(BATCH_SIZE, 100)) _ = s.run( d_trainer, feed_dict={real_images: batch_images, seed: batch_seed} ) _ = s.run(g_trainer, feed_dict={seed: batch_seed}) print("on epoch{}".format(epoch)) sample_seed = np.random.uniform(-1, 1, size=(1, 100)) gen_sample = s.run(generator(seed, reuse=True), feed_dict={seed: sample_seed}) samples.append(gen_sample) plt.imshow(samples[0].reshape(28, 28)) plt.imshow(samples[99].reshape(28, 28))