parser = argparse.ArgumentParser(description="DeepLab-ResNet Network")

parser.add_argument("--batch-size", type=int, default=BATCH_SIZE,

help="Number of images sent to the network in one step.")

# parser.add_argument("--data-dir", type=str, default=DATA_DIRECTORY,

# help="Path to the directory containing the PASCAL VOC dataset.")

parser.add_argument("--data-list", type=str, default=DATA_LIST_PATH,

help="Path to the file listing the images in the dataset.")

parser.add_argument("--ignore-label", type=int, default=IGNORE_LABEL,

help="The index of the label to ignore during the training.")

parser.add_argument("--input-size", type=str, default=INPUT_SIZE,

help="Comma-separated string with height and width of images.")

parser.add_argument("--is-training", action="store_true",

help="Whether to updates the running means and variances during the training.")

parser.add_argument("--learning-rate", type=float, default=LEARNING_RATE,

help="Base learning rate for training with polynomial decay.")

parser.add_argument("--momentum", type=float, default=MOMENTUM,

help="Momentum component of the optimiser.")

parser.add_argument("--not-restore-last", action="store_true",

help="Whether to not restore last (FC) layers.")

parser.add_argument("--num-classes", type=int, default=NUM_CLASSES,

help="Number of classes to predict (including background).")

parser.add_argument("--num-steps", type=int, default=NUM_STEPS,

help="Number of training steps.")

parser.add_argument("--power", type=float, default=POWER,

help="Decay parameter to compute the learning rate.")

parser.add_argument("--random-mirror", action="store_true",

help="Whether to randomly mirror the inputs during the training.")

parser.add_argument("--random-scale", action="store_true",

help="Whether to randomly scale the inputs during the training.")

parser.add_argument("--random-seed", type=int, default=RANDOM_SEED,

help="Random seed to have reproducible results.")

parser.add_argument("--restore-from", type=str, default=RESTORE_FROM,

help="Where restore model parameters from.")

parser.add_argument("--save-num-images", type=int, default=SAVE_NUM_IMAGES,

help="How many images to save.")

parser.add_argument("--save-pred-every", type=int, default=SAVE_PRED_EVERY,

help="Save summaries and checkpoint every often.")

parser.add_argument("--snapshot-dir", type=str, default=SNAPSHOT_DIR,

help="Where to save snapshots of the model.")

parser.add_argument("--weight-decay", type=float, default=WEIGHT_DECAY,

help="Regularisation parameter for L2-loss.")

parser.add_argument("--update-mean-var", action="store_true",

help="whether to get update_op from tf.Graphic_Keys")

model_name = 'model.ckpt'

checkpoint_path = os.path.join(logdir, model_name)

if not os.path.exists(logdir):

os.makedirs(logdir)

saver.save(sess, checkpoint_path, global_step=step)

"""Create the model and start the training."""

args = get_arguments()

h, w = map(int, args.input_size.split(','))

input_size = (h, w)

tf.set_random_seed(args.random_seed)

coord = tf.train.Coordinator()

with tf.name_scope("create_inputs"):

reader = ImageReader(

args.data_list,

input_size,

args.random_scale,

args.random_mirror,

args.ignore_label,

IMG_MEAN,

coord)

image_batch, label_batch = reader.dequeue(args.batch_size)

net = PSPNet50({'data': image_batch}, is_training=True, num_classes=args.num_classes)

raw_output = net.layers['conv6']

# According from the prototxt in Caffe implement, learning rate must multiply by 10.0 in pyramid module

fc_list = ['conv5_3_pool1_conv', 'conv5_3_pool2_conv', 'conv5_3_pool3_conv', 'conv5_3_pool6_conv', 'conv6', 'conv5_4']

restore_var = [v for v in tf.global_variables() if not (len([f for f in fc_list if f in v.name])) or not args.not_restore_last]

all_trainable = [v for v in tf.trainable_variables() if 'gamma' not in v.name and 'beta' not in v.name]

fc_trainable = [v for v in all_trainable if v.name.split('/')[0] in fc_list]

conv_trainable = [v for v in all_trainable if v.name.split('/')[0] not in fc_list] # lr * 1.0

fc_w_trainable = [v for v in fc_trainable if 'weights' in v.name] # lr * 10.0

fc_b_trainable = [v for v in fc_trainable if 'biases' in v.name] # lr * 20.0

assert(len(all_trainable) == len(fc_trainable) + len(conv_trainable))

assert(len(fc_trainable) == len(fc_w_trainable) + len(fc_b_trainable))

# Predictions: ignoring all predictions with labels greater or equal than n_classes

raw_prediction = tf.reshape(raw_output, [-1, args.num_classes])

label_proc = prepare_label(label_batch, tf.stack(raw_output.get_shape()[1:3]), num_classes=args.num_classes, one_hot=False) # [batch_size, h, w]

raw_gt = tf.reshape(label_proc, [-1,])

indices = tf.squeeze(tf.where(tf.less_equal(raw_gt, args.num_classes - 1)), 1)

gt = tf.cast(tf.gather(raw_gt, indices), tf.int32)

prediction = tf.gather(raw_prediction, indices)

# Pixel-wise softmax loss.

loss = tf.nn.sparse_softmax_cross_entropy_with_logits(logits = prediction, labels = gt)

# Make mistakes for class N more important for network

if USE_CLASS_WEIGHTS:

if len(CLASS_WEIGHTS) != NUM_CLASSES:

print('Incorrect class weights, it will be not used')

else:

mask = tf.zeros_like(loss)

for i, w in enumerate(CLASS_WEIGHTS):

mask = mask + tf.cast(tf.equal(gt, i), tf.float32) * tf.constant(w)

loss = loss * mask

l2_losses = [args.weight_decay * tf.nn.l2_loss(v) for v in tf.trainable_variables() if 'weights' in v.name]

reduced_loss = tf.reduce_mean(loss) + tf.add_n(l2_losses)

# Processed predictions: for visualisation.

raw_output_up = tf.image.resize_bilinear(raw_output, tf.shape(image_batch)[1:3,])

raw_output_up = tf.argmax(raw_output_up, dimension=3)

pred = tf.expand_dims(raw_output_up, dim=3)

# Image summary.

images_summary = tf.py_func(inv_preprocess, [image_batch, args.save_num_images, IMG_MEAN], tf.uint8)

labels_summary = tf.py_func(decode_labels, [label_batch, args.save_num_images, args.num_classes], tf.uint8)

preds_summary = tf.py_func(decode_labels, [pred, args.save_num_images, args.num_classes], tf.uint8)

total_summary = tf.summary.image('images',

tf.concat(axis=2, values=[images_summary, labels_summary, preds_summary]),

max_outputs=args.save_num_images) # Concatenate row-wise.

summary_writer = tf.summary.FileWriter(args.snapshot_dir,

graph=tf.get_default_graph())

# Using Poly learning rate policy

base_lr = tf.constant(args.learning_rate)

step_ph = tf.placeholder(dtype=tf.float32, shape=())

learning_rate = tf.scalar_mul(base_lr, tf.pow((1 - step_ph / args.num_steps), args.power))

# Gets moving_mean and moving_variance update operations from tf.GraphKeys.UPDATE_OPS

if args.update_mean_var == False:

update_ops = None

else:

update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)

with tf.control_dependencies(update_ops):

opt_conv = tf.train.MomentumOptimizer(learning_rate, args.momentum)

opt_fc_w = tf.train.MomentumOptimizer(learning_rate * 10.0, args.momentum)

opt_fc_b = tf.train.MomentumOptimizer(learning_rate * 20.0, args.momentum)

grads = tf.gradients(reduced_loss, conv_trainable + fc_w_trainable + fc_b_trainable)

grads_conv = grads[:len(conv_trainable)]

grads_fc_w = grads[len(conv_trainable) : (len(conv_trainable) + len(fc_w_trainable))]

grads_fc_b = grads[(len(conv_trainable) + len(fc_w_trainable)):]

train_op_conv = opt_conv.apply_gradients(zip(grads_conv, conv_trainable))

train_op_fc_w = opt_fc_w.apply_gradients(zip(grads_fc_w, fc_w_trainable))

train_op_fc_b = opt_fc_b.apply_gradients(zip(grads_fc_b, fc_b_trainable))

train_op = tf.group(train_op_conv, train_op_fc_w, train_op_fc_b)

# Set up tf session and initialize variables.

config = tf.ConfigProto()

# config.gpu_options.allow_growth = True

# config.allow_soft_placement = True

# config.intra_op_parallelism_threads = 1

sess = tf.Session(config=config)

init = tf.global_variables_initializer()

sess.run(init)

# Saver for storing checkpoints of the model.

saver = tf.train.Saver(var_list=tf.global_variables(), max_to_keep=10)

ckpt = tf.train.get_checkpoint_state(SNAPSHOT_DIR)

if ckpt and ckpt.model_checkpoint_path:

loader = tf.train.Saver(var_list=restore_var)

load_step = int(os.path.basename(ckpt.model_checkpoint_path).split('-')[1])

load(loader, sess, ckpt.model_checkpoint_path)

else:

print('No checkpoint file found.')

load_step = 0

# Start queue threads.

threads = tf.train.start_queue_runners(coord=coord, sess=sess)

# Iterate over training steps.

for step in range(args.num_steps):

start_time = time.time()

feed_dict = {step_ph: step}

if step % args.save_pred_every == 0:

loss_value, _, summary = sess.run([reduced_loss, train_op, total_summary], feed_dict=feed_dict)

summary_writer.add_summary(summary, step)

save(saver, sess, args.snapshot_dir, step)

else:

z, t, o, p, loss_value, _ = sess.run([raw_gt, raw_output, gt, prediction, reduced_loss, train_op], feed_dict=feed_dict)

print(z.shape, t.shape, o.shape, p.shape)

duration = time.time() - start_time

print('step {:d} \t loss = {:.3f}, ({:.3f} sec/step)'.format(step, loss_value, duration))

coord.request_stop()