| 36 | |
| 37 | |
| 38 | def train(args): |
| 39 | |
| 40 | |
| 41 | input_photo = tf.placeholder(tf.float32, [args.batch_size, |
| 42 | args.patch_size, args.patch_size, 3]) |
| 43 | |
| 44 | output = network.unet_generator(input_photo) |
| 45 | |
| 46 | recon_loss = tf.reduce_mean(tf.losses.absolute_difference(input_photo, output)) |
| 47 | |
| 48 | all_vars = tf.trainable_variables() |
| 49 | gene_vars = [var for var in all_vars if 'gene' in var.name] |
| 50 | |
| 51 | update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) |
| 52 | with tf.control_dependencies(update_ops): |
| 53 | |
| 54 | optim = tf.train.AdamOptimizer(args.adv_train_lr, beta1=0.5, beta2=0.99)\ |
| 55 | .minimize(recon_loss, var_list=gene_vars) |
| 56 | |
| 57 | |
| 58 | ''' |
| 59 | config = tf.ConfigProto() |
| 60 | config.gpu_options.allow_growth = True |
| 61 | sess = tf.Session(config=config) |
| 62 | ''' |
| 63 | gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=args.gpu_fraction) |
| 64 | sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options)) |
| 65 | saver = tf.train.Saver(var_list=gene_vars, max_to_keep=20) |
| 66 | |
| 67 | with tf.device('/device:GPU:0'): |
| 68 | |
| 69 | sess.run(tf.global_variables_initializer()) |
| 70 | face_photo_dir = 'dataset/photo_face' |
| 71 | face_photo_list = utils.load_image_list(face_photo_dir) |
| 72 | scenery_photo_dir = 'dataset/photo_scenery' |
| 73 | scenery_photo_list = utils.load_image_list(scenery_photo_dir) |
| 74 | |
| 75 | |
| 76 | for total_iter in tqdm(range(args.total_iter)): |
| 77 | |
| 78 | if np.mod(total_iter, 5) == 0: |
| 79 | photo_batch = utils.next_batch(face_photo_list, args.batch_size) |
| 80 | else: |
| 81 | photo_batch = utils.next_batch(scenery_photo_list, args.batch_size) |
| 82 | |
| 83 | _, r_loss = sess.run([optim, recon_loss], feed_dict={input_photo: photo_batch}) |
| 84 | |
| 85 | if np.mod(total_iter+1, 50) == 0: |
| 86 | |
| 87 | print('pretrain, iter: {}, recon_loss: {}'.format(total_iter, r_loss)) |
| 88 | if np.mod(total_iter+1, 500 ) == 0: |
| 89 | saver.save(sess, args.save_dir+'save_models/model', |
| 90 | write_meta_graph=False, global_step=total_iter) |
| 91 | |
| 92 | photo_face = utils.next_batch(face_photo_list, args.batch_size) |
| 93 | photo_scenery = utils.next_batch(scenery_photo_list, args.batch_size) |
| 94 | |
| 95 | result_face = sess.run(output, feed_dict={input_photo: photo_face}) |