def main(argv = None):
"""
tensorflow 中的app.run会先解析命令行参数flag，然后执行main函数
"""
# 定义存放图像，标签和过拟合的数据结构
keep_probability = tf.placeholder(tf.float32, name="keep_probabilty")
image = tf.placeholder(tf.float32, shape=[IMAGE_SIZE, IMAGE_SIZE, 3], name="input_image")
annotation = tf.placeholder(tf.int32, shape=[IMAGE_SIZE, IMAGE_SIZE, 1], name="annotation")
pred_annotation, logits = segmentation(image, keep_probability)
# 添加监控信息，可以通过tensorboard查看
tf.summary.image("input_image", image, max_outputs = 2)
tf.summary.image("ground_truth", tf.cast(annotation, tf.uint8), max_outputs = 2)
tf.summary.image("pred_annotation", tf.cast(pred_annotation, tf.uint8), max_outputs = 2)
# 还没看懂loss函数是怎么计算的，reduce_mean是求平均值，squeeze为去掉维度是1
loss = tf.reduce_mean((tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits,
labels=tf.squeeze(annotation, squeeze_dims=[3]),
name="entropy")))
tf.summary.scalar("entropy", loss)
# 获取在训练中的变量列表
trainable_var = tf.trainable_variables()
if FLAGS.debug:
for var in trainable_var:
utils.add_to_regularization_and_summary(var)
train_op = train(loss, trainable_var)
print("Setting up summary op...")
summary_op = tf.summary.merge_all()
print("Setting up image reader...")
sess = tf.Session()
print("Setting up Saver...")
saver = tf.train.Saver()
summary_writer = tf.summary.FileWriter(FLAGS.logs_dir, sess.graph)
sess.run(tf.global_variables_initializer())
name_list = shuffle_namelist(train_Path)
start_index = 0
batch_size = FLAGS.batch_size
if FLAGS.mode == "train":
for itr in range(MAX_ITERATION):
train_images, train_annotations, start_index = read_2_namelist(name_list, batch_size, start_index)
feed_dict = {image: train_images, annotation: train_annotations, keep_probability: 0}
sess.run(train_op, feed_dict=feed_dict)
if itr % 10 == 0:
train_loss, summary_str = sess.run([loss, summary_op], feed_dict=feed_dict)
print("Step: %d, Train_loss:%g" % (itr, train_loss))
summary_writer.add_summary(summary_str, itr)
if itr % 500 == 0:
saver.save(sess, FLAGS.logs_dir + "FCN_model.ckpt", itr)
if FLAGS.mode == "predict":
saver.restore(sess, FLAGS.logs_dir + "FCN_model.ckpt")
predict(sess, pred_annotation)
def create_network(self, generator_dims, discriminator_dims, optimizer="Adam", learning_rate=2e-4,
optimizer_param=0.9, improved_gan_loss=True):
print("Setting up model...")
self._setup_placeholder()
tf.summary.histogram("z", self.z_vec)
self.gen_images = self._generator(self.z_vec, generator_dims, self.train_phase, scope_name="generator")
tf.summary.image("image_real", self.images, max_outputs=2)
tf.summary.image("image_generated", self.gen_images, max_outputs=2)
def leaky_relu(x, name="leaky_relu"):
return utils.leaky_relu(x, alpha=0.2, name=name)
discriminator_real_prob, logits_real, feature_real = self._discriminator(self.images, discriminator_dims,
self.train_phase,
activation=leaky_relu,
scope_name="discriminator",
scope_reuse=False)
discriminator_fake_prob, logits_fake, feature_fake = self._discriminator(self.gen_images, discriminator_dims,
self.train_phase,
activation=leaky_relu,
scope_name="discriminator",
scope_reuse=True)
# utils.add_activation_summary(tf.identity(discriminator_real_prob, name='disc_real_prob'))
# utils.add_activation_summary(tf.identity(discriminator_fake_prob, name='disc_fake_prob'))
# Loss calculation
self._gan_loss(logits_real, logits_fake, feature_real, feature_fake, use_features=improved_gan_loss)
train_variables = tf.trainable_variables()
for v in train_variables:
# print (v.op.name)
utils.add_to_regularization_and_summary(var=v)
self.generator_variables = [v for v in train_variables if v.name.startswith("generator")]
# print(map(lambda x: x.op.name, generator_variables))
self.discriminator_variables = [v for v in train_variables if v.name.startswith("discriminator")]
# print(map(lambda x: x.op.name, discriminator_variables))
optim = self._get_optimizer(optimizer, learning_rate, optimizer_param)
self.generator_train_op = self._train(self.gen_loss, self.generator_variables, optim)
self.discriminator_train_op = self._train(self.discriminator_loss, self.discriminator_variables, optim)
def create_network(self,
generator_dims,
discriminator_dims,
optimizer="Adam",
learning_rate=2e-4,
optimizer_param=0.9,
improved_gan_loss=True,
trainable_z=False,
trainable_image=False):
print("Setting up model...")
self._setup_placeholder()
tf.summary.histogram("z", self.z_vec)
if trainable_z:
# make z iterator variable
self.z_iterator = tf.Variable(np.random.uniform(
-1.0, 1.0, (self.batch_size, int(
self.z_vec.get_shape()[1]))).astype(dtype=np.float32),
name="z_iterator")
self.init_z_iterator = tf.group(self.z_iterator.assign(self.z_vec))
self.z_iterator_min_max = tf.maximum(
tf.minimum(tf.maximum(self.z_iterator, -1.0), 1.0), -1.0)
self.z_vec_in = self.z_iterator_min_max
else:
self.z_vec_in = self.z_vec
# generator for training
self.gen_images = self._generator(self.z_vec_in,
generator_dims,
self.train_phase,
scope_name="generator")
if trainable_image:
# make image iterator variable
self.image_iterator = tf.Variable(np.random.uniform(
0.0, 1.0,
(self.batch_size, 64, 64, 3)).astype(dtype=np.float32),
name="image_iterator")
self.init_image_iterator = tf.group(
self.image_iterator.assign(self.gen_images))
self.image_iterator_min_max = tf.minimum(
tf.maximum(self.image_iterator, -1.0), 1.0)
self.gen_images_out = self.image_iterator_min_max
else:
self.gen_images_out = self.gen_images
# generator for z iterator
tf.summary.image("image_real", self.images)
tf.summary.image("image_generated", self.gen_images_out)
def leaky_relu(x, name="leaky_relu"):
return utils.leaky_relu(x, alpha=0.2, name=name)
discriminator_real_prob, logits_real, feature_real = self._discriminator(
self.images,
discriminator_dims,
self.train_phase,
activation=leaky_relu,
scope_name="discriminator",
scope_reuse=False)
discriminator_fake_prob, logits_fake, feature_fake = self._discriminator(
self.gen_images_out,
discriminator_dims,
self.train_phase,
activation=leaky_relu,
scope_name="discriminator",
scope_reuse=True)
# utils.add_activation_summary(tf.identity(discriminator_real_prob, name='disc_real_prob'))
# utils.add_activation_summary(tf.identity(discriminator_fake_prob, name='disc_fake_prob'))
# Loss calculation
self._gan_loss(logits_real,
logits_fake,
feature_real,
feature_fake,
use_features=improved_gan_loss)
train_variables = tf.trainable_variables()
for v in train_variables:
utils.add_to_regularization_and_summary(var=v)
# get variable lists for everything
self.generator_variables = [
v for v in train_variables if v.name.startswith("generator")
]
self.discriminator_variables = [
v for v in train_variables if v.name.startswith("discriminator")
]
self.image_iterator_variables = [
v for v in train_variables if v.name.startswith("image_iterator")
]
self.z_iterator_variables = [
v for v in train_variables if v.name.startswith("z_iterator")
]
# set optimizer
optim = self._get_optimizer(optimizer, learning_rate, optimizer_param)
optim_z = self._get_optimizer(optimizer, learning_rate * 10.0,
optimizer_param)
# make train ops
if not trainable_image and not trainable_z:
self.generator_train_op = self._train(self.gen_loss,
self.generator_variables,
optim)
self.discriminator_train_op = self._train(
self.discriminator_loss, self.discriminator_variables, optim)
if trainable_image:
self.image_iterator_train_op = self._train(
self.gen_loss, self.image_iterator_variables, optim)
if trainable_z:
self.z_iterator_train_op = self._train(self.gen_loss,
self.z_iterator_variables,
optim_z)