def main(): """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) # Create queue coordinator. coord = tf.train.Coordinator() # Load reader. with tf.name_scope("create_inputs"): reader = ImageReader(args.data_dir, 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) # Create network. sys.stdout.flush() sys.stderr.flush() net = DeepLabResNetModel({'data': image_batch}, is_training=args.is_training, num_classes=args.num_classes) sys.stdout.flush() sys.stderr.flush() # For a small batch size, it is better to keep # the statistics of the BN layers (running means and variances) # frozen, and to not update the values provided by the pre-trained model. # If is_training=True, the statistics will be updated during the training. # Note that is_training=False still updates BN parameters gamma (scale) and beta (offset) # if they are presented in var_list of the optimiser definition. # Predictions. raw_output = net.layers['fc1_voc12'] # Which variables to load. Running means and variances are not trainable, # thus all_variables() should be restored. all_trainable = [ v for v in tf.trainable_variables() if 'beta' not in v.name and 'gamma' not in v.name ] fc_trainable = [v for v in all_trainable if 'fc' in v.name] conv_trainable = [v for v in all_trainable if 'fc' not in v.name] # 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 if args.freeze_convolutions: print("freezing backbone") all_trainable = [v for v in all_trainable if v not in conv_trainable] conv_trainable = [] assert (len(all_trainable) == len(fc_trainable) + len(conv_trainable)) assert (len(fc_trainable) == len(fc_w_trainable) + len(fc_b_trainable)) print("training {} variables out of total {} trainable variables".format( len(all_trainable), len(tf.trainable_variables()))) # 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) 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) if not os.path.exists(args.snapshot_dir): os.mkdir(args.snapshot_dir) 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()) # Define loss and optimisation parameters. 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)) 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_group_list = [] if len(conv_trainable) > 0: train_op_conv = opt_conv.apply_gradients( zip(grads_conv, conv_trainable)) train_op_group_list.append(train_op_conv) train_op_fc_w = opt_fc_w.apply_gradients(zip(grads_fc_w, fc_w_trainable)) train_op_group_list.append(train_op_fc_w) train_op_fc_b = opt_fc_b.apply_gradients(zip(grads_fc_b, fc_b_trainable)) train_op_group_list.append(train_op_fc_b) train_op = tf.group(*train_op_group_list) # Set up tf session and initialize variables. config = tf.ConfigProto() config.gpu_options.allow_growth = True sess = tf.Session(config=config) init = tf.global_variables_initializer() sess.run(init) snapshot_dir = args.snapshot_dir skip_next_arg_print = False for i in range(1, len(sys.argv)): if skip_next_arg_print: skip_next_arg_print = False continue if sys.argv[i].startswith('--snapshot-dir') or sys.argv[i].startswith( '--data-dir') or sys.argv[i].startswith('--restore-from'): skip_next_arg_print = True continue if sys.argv[i].startswith('--'): words = sys.argv[i].replace('--', '').split('-') snapshot_dir += '_' + ''.join([w[0] for w in words]) + "=" else: snapshot_dir += sys.argv[i].replace('/', '-') # Saver for storing checkpoints of the model. saver = tf.train.Saver(var_list=tf.global_variables(), max_to_keep=50) start_from_step = 0 # Load variables if the checkpoint is provided. if args.restore_from is not None: if args.restore_from == RESTORE_FROM and os.path.exists( snapshot_dir) and len(os.listdir(snapshot_dir)) >= 3: restore_path = tf.train.latest_checkpoint(snapshot_dir) print("Auto restoring weights from " + str(restore_path)) else: restore_path = args.restore_from try: start_from_step = int(restore_path.split("-")[-1]) except ValueError: start_from_step = 0 print("Auto starting from step " + str(start_from_step) + " (detected from checkpoint file)") vars_in_checkpoint = get_tensors_in_checkpoint_file( file_name=restore_path) loadable_tensors = match_loaded_and_memory_tensors(vars_in_checkpoint) loadable_tensors = [ v for v in loadable_tensors if 'fc' not in v.name or not args.not_restore_last ] loader = tf.train.Saver(var_list=loadable_tensors) load(loader, sess, restore_path) sys.stdout.flush() sys.stderr.flush() # Start queue threads. threads = tf.train.start_queue_runners(coord=coord, sess=sess) os.makedirs(snapshot_dir, exist_ok=True) with open(os.path.join(snapshot_dir, "args.json"), "+w") as f: f.write(json.dumps(args.__dict__, indent=2)) def should_print(step): if step < 100: return True elif step < 1000: return step % 100 == 0 elif step < 10000: return step % 1000 == 0 else: return step % 10000 == 0 loss_sum = 0 num_loss_sum = 0 # Iterate over training steps. for step in range(start_from_step, args.num_steps): start_time = time.time() feed_dict = {step_ph: step} if step % args.save_pred_every == 0 or step == args.num_steps - 1: loss_value, imgs, lbls, preds, summary, _ = sess.run( [ reduced_loss, image_batch, label_batch, pred, total_summary, train_op ], feed_dict=feed_dict) summary_writer.add_summary(summary, step) save(saver, sess, snapshot_dir, step) else: loss_value, _ = sess.run([reduced_loss, train_op], feed_dict=feed_dict) loss_sum += loss_value num_loss_sum += 1 duration = time.time() - start_time if should_print(step) or should_print( step - start_from_step) or step == args.num_steps - 1: print( '{:2.2f}% step {:d}/{:d} \t loss = {:.3f} , ({:.3f} sec/step)'. format( float(step / args.num_steps) * 100., step, args.num_steps, loss_sum / num_loss_sum, duration)) loss_sum = 0 num_loss_sum = 0 sys.stdout.flush() sys.stderr.flush() coord.request_stop() coord.join(threads)
def main(): """Create the model and start the training.""" args = get_arguments() # Load reader. h, w = map(int, cfg.INPUT_SIZE.split(',')) c = 4 if cfg.ONLY_POS else 5 reader_option = {"resize":True, "resize_size":[h,w]} train_dataset_reader = BatchDataset(reader_option) num_file = train_dataset_reader.get_image_number() num_steps = args.num_epochs * num_file # Create network. os.environ['CUDA_VISIBLE_DEVICES'] = '0' image_batch = tf.placeholder(tf.float32, shape=[args.batch_size, h, w, c], name='input') label_batch = tf.placeholder(tf.uint8, shape=[args.batch_size, h, w, 1], name='label') net = DeepLabResNetModel({'data': image_batch}, is_training=args.is_training) # For a small batch size, it is better to keep # the statistics of the BN layers (running means and variances) # frozen, and to not update the values provided by the pre-trained model. # If is_training=True, the statistics will be updated during the training. # Note that is_training=False still updates BN parameters gamma (scale) and beta (offset) # if they are presented in var_list of the optimiser definition. # Predictions. raw_output = net.layers['fc1_voc12'] # Which variables to load. Running means and variances are not trainable, # thus all_variables() should be restored. restore_var = tf.global_variables() all_trainable = [v for v in tf.trainable_variables() if 'beta' not in v.name and 'gamma' not in v.name] fc_trainable = [v for v in all_trainable if 'fc' in v.name] # lr * 10.0 conv1_trainable = [v for v in all_trainable if 'conv1' in v.name] # lr * 20.0 conv_trainable = [v for v in all_trainable if 'fc' not in v.name and 'conv1' not in v.name] # lr * 1.0 assert(len(all_trainable) == len(fc_trainable) + len(conv1_trainable) + len(conv_trainable)) # Predictions: ignoring all predictions with labels greater or equal than n_classes raw_prediction = tf.reshape(raw_output, [-1, cfg.num_classes]) label_proc = prepare_label(label_batch, tf.stack(raw_output.get_shape()[1:3]), one_hot=False) # [batch_size, h, w] raw_gt = tf.reshape(label_proc, [-1,]) raw_prediction = tf.reshape(raw_prediction, [-1,]) gt = tf.cast(raw_gt, tf.float32) prediction = tf.cast(raw_prediction, tf.float32) # Pixel-wise softmax loss. # loss = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=prediction, labels=gt) loss = tf.nn.weighted_cross_entropy_with_logits(targets=gt, logits=prediction, pos_weight=5, name='weighted_sigmoid') 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)*20 + 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], tf.uint8) labels_summary = tf.py_func(decode_labels, [label_batch, args.save_num_images], tf.uint8) preds_summary = tf.py_func(decode_labels, [pred, args.save_num_images], tf.uint8) total_summary = tf.summary.image('images', tf.concat([images_summary, labels_summary, preds_summary], axis=2), max_outputs=args.save_num_images) # Concatenate row-wise. summary_writer = tf.summary.FileWriter(args.snapshot_dir) ''' # Define loss and optimisation parameters. base_lr = tf.constant(args.learning_rate) step_ph = tf.placeholder(dtype=tf.float32, shape=()) learning_rate = tf.scalar_mul(base_lr, tf.pow(args.power, step_ph // num_steps)) opt_conv = tf.train.MomentumOptimizer(learning_rate*1, args.momentum) opt_fc = tf.train.MomentumOptimizer(learning_rate*5.0, args.momentum) opt_conv1 = tf.train.MomentumOptimizer(learning_rate*5.0, args.momentum) grads = tf.gradients(reduced_loss, conv_trainable + fc_trainable + conv1_trainable) grads_conv = grads[:len(conv_trainable)] grads_fc = grads[len(conv_trainable) : (len(conv_trainable) + len(fc_trainable))] grads_conv1 = grads[(len(conv_trainable) + len(fc_trainable)):] train_op_conv = opt_conv.apply_gradients(zip(grads_conv, conv_trainable)) train_op_fc = opt_fc.apply_gradients(zip(grads_fc, fc_trainable)) train_op_conv1 = opt_conv1.apply_gradients(zip(grads_conv1, conv1_trainable)) train_op = tf.group(train_op_conv, train_op_fc, train_op_conv1) # evaluation # yjl::check the update operation pred = tf.argmax(raw_predictions, axis=-1) mIoU, update_op = tf.contrib.metrics.streaming_mean_iou(pred, raw_gt, num_classes=n_classes, weights=weights) # Set up tf session and initialize variables. config = tf.ConfigProto(allow_soft_placement=True) config.gpu_options.allow_growth = True 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=restore_var, max_to_keep=2)#keep_checkpoint_every_n_hours=1.0) # Load variables if the checkpoint is provided. # load_var_list = [v for v in restore_var if ('conv1' not in v.name) and ('fc1_voc12' not in v.name)] if args.restore_from is not None: loader = tf.train.Saver(var_list=restore_var) # loader = tf.train.Saver(var_list=load_var_list) load(loader, sess, args.restore_from) pdb.set_trace() for step in range(num_steps): start_time = time.time() images, labels = train_dataset_reader.next_batch(args.batch_size) feed_dict = {image_batch:images, label_batch:labels, step_ph:step} #feed_dict = {step_ph:step} if step % args.save_pred_every == 0: # loss_value, preds, summary = sess.run([reduced_loss, pred, total_summary], feed_dict=feed_dict) # summary_writer.add_summary(summary, step) loss_value = sess.run([reduced_loss], feed_dict=feed_dict) save(saver, sess, args.snapshot_dir, step) loss_value, inf_loss, l2_loss, _ = sess.run([reduced_loss, loss, l2_losses, train_op], feed_dict=feed_dict) duration = time.time() - start_time epoch = (int)(step*args.batch_size/num_file) print('epoch {:d} /step {:d} \t loss = {:.3f}, ({:.3f} sec/step)'.format(epoch, step, loss_value, duration))
def main(): """Create the model and start the training.""" args = get_arguments() print(args) if args.not_restore_last: try: shutil.rmtree(args.snapshot_dir) except Exception as e: print(e) os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_mask h, w = map(int, args.input_size.split(',')) input_size = (h, w) tf.set_random_seed(args.random_seed) # Create queue coordinator. coord = tf.train.Coordinator() # Load reader. mode = tf.placeholder(tf.bool, shape=()) with tf.name_scope("create_inputs"): reader = ImageReader( args.data_dir, args.data_list, input_size, args.random_scale, args.random_mirror, args.ignore_label, IMG_MEAN, coord) image_batch_train, label_batch_train = reader.dequeue(args.batch_size) with tf.name_scope("val_inputs"): reader = ImageReader( args.data_dir, args.val_data_list, input_size, args.random_scale, args.random_mirror, args.ignore_label, IMG_MEAN, coord) image_batch_val, label_batch_val = reader.dequeue(args.batch_size) image_batch = tf.cond(mode, lambda: image_batch_train, lambda: image_batch_val) label_batch = tf.cond(mode, lambda: label_batch_train, lambda: label_batch_val) # Create network. net = DeepLabResNetModel({'data': image_batch}, is_training=args.is_training, num_classes=args.num_classes) # For a small batch size, it is better to keep # the statistics of the BN layers (running means and variances) # frozen, and to not update the values provided by the pre-trained model. # If is_training=True, the statistics will be updated during the training. # Note that is_training=False still updates BN parameters gamma (scale) and beta (offset) # if they are presented in var_list of the optimiser definition. # Predictions. raw_output = net.layers['fc1_voc12'] # Which variables to load. Running means and variances are not trainable, # thus all_variables() should be restored. restore_var = [v for v in tf.global_variables() if 'fc' not in v.name or not args.not_restore_last] all_trainable = [v for v in tf.trainable_variables() if 'beta' not in v.name and 'gamma' not in v.name] fc_trainable = [v for v in all_trainable if 'fc' in v.name] conv_trainable = [v for v in all_trainable if 'fc' not in v.name] # 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) output_op = tf.cast(tf.argmax(prediction, axis=-1), tf.int32) correct_pred = tf.equal(output_op, gt) accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32)) # Pixel-wise softmax loss. loss = [] accuracy_per_class = [] softmax_weights_per_class = tf.constant(LUNA16_softmax_weights, dtype=tf.float32) for i in xrange(0, args.num_classes): curr_class = tf.constant(i, tf.int32) loss.append(softmax_weights_per_class[i] * tf.losses.sparse_softmax_cross_entropy(logits=prediction, labels=gt, weights=tf.where( tf.equal(gt, curr_class), tf.zeros_like(gt), tf.ones_like(gt)))) accuracy_per_class.append( tf.reduce_mean(tf.cast(tf.gather(correct_pred, tf.where(tf.equal(gt, curr_class))), tf.float32))) 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(tf.stack(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. reduced_loss_train = tf.Variable(0, trainable=False, dtype=tf.float32) accuracy_train = tf.Variable(0, trainable=False, dtype=tf.float32) reduced_loss_val = tf.Variable(0, trainable=False, dtype=tf.float32) accuracy_val = tf.Variable(0, trainable=False, dtype=tf.float32) reduced_loss_train = tf.cond(mode, lambda: tf.assign(reduced_loss_train, reduced_loss), lambda: reduced_loss_train) accuracy_train = tf.cond(mode, lambda: tf.assign(accuracy_train, accuracy), lambda: accuracy_train) reduced_loss_val = tf.cond(mode, lambda: reduced_loss_val, lambda: tf.assign(reduced_loss_val, reduced_loss)) accuracy_val = tf.cond(mode, lambda: accuracy_val, lambda: tf.assign(accuracy_val, accuracy)) accuracy_per_class_train = [] accuracy_per_class_val = [] for i in xrange(0, args.num_classes): temp_train_var = tf.Variable(0, trainable=False, dtype=tf.float32) temp_val_var = tf.Variable(0, trainable=False, dtype=tf.float32) accuracy_per_class_train.append( tf.cond(mode, lambda: tf.assign(temp_train_var, accuracy_per_class[i]), lambda: temp_train_var)) accuracy_per_class_val.append( tf.cond(mode, lambda: temp_val_var, lambda: tf.assign(temp_val_var, accuracy_per_class[i]))) accuracy_output = tf.cond(mode, lambda: accuracy_train, lambda: accuracy_val) loss_output = tf.cond(mode, lambda: reduced_loss_train, lambda: reduced_loss_val) tf.summary.scalar("Loss", loss_output, collections=['all']) tf.summary.scalar("Accuracy", accuracy_output, collections=['all']) 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) counter_no_reset = tf.Variable(tf.zeros([2, args.num_classes]), trainable=False, dtype=tf.float32) counter = tf.Variable(tf.zeros([2, args.num_classes]), trainable=False, dtype=tf.float32) counter_no_reset_val = tf.Variable(tf.zeros([2, args.num_classes]), trainable=False, dtype=tf.float32) counter_val = tf.Variable(tf.zeros([2, args.num_classes]), trainable=False, dtype=tf.float32) step_ph = tf.placeholder(dtype=tf.float32, shape=()) counter, counter_no_reset = tf.cond(mode, lambda: tf.py_func(update_IoU, [tf.squeeze(pred, axis=-1), tf.squeeze(label_batch, axis=-1), counter, counter_no_reset, args.num_classes, args.batch_size, step_ph, args.save_pred_every], [tf.float32, tf.float32]), lambda: [counter, counter_no_reset]) counter_val, counter_no_reset_val = tf.cond(mode, lambda: [counter_val, counter_no_reset_val], lambda: tf.py_func(update_IoU, [tf.squeeze(pred, axis=-1), tf.squeeze(label_batch, axis=-1), counter_val, counter_no_reset_val, args.num_classes, args.batch_size, step_ph, args.save_pred_every], [tf.float32, tf.float32])) eps = tf.constant(1e-10, dtype=tf.float32) IoU_summary = counter[0] / tf.add(eps, counter[1]) IoU_summary_no_reset = counter_no_reset[0] / tf.add(eps, counter_no_reset[1]) Val_IoU_summary = counter_val[0] / tf.add(eps, counter_val[1]) Val_IoU_summary_no_reset = counter_no_reset_val[0] / tf.add(eps, counter_no_reset_val[1]) mIoU = tf.reduce_mean(IoU_summary) mIoU_no_reset = tf.reduce_mean(IoU_summary_no_reset) Val_mIoU = tf.reduce_mean(Val_IoU_summary) Val_mIoU_no_reset = tf.reduce_mean(Val_IoU_summary_no_reset) IoU_summary_output_intermed = tf.cond(mode, lambda: IoU_summary, lambda: Val_IoU_summary) IoU_summary_no_reset_output_intermed = tf.cond(mode, lambda: IoU_summary_no_reset, lambda: Val_IoU_summary_no_reset) accuracy_per_class_output_intermed = tf.cond(mode, lambda: accuracy_per_class_train, lambda: accuracy_per_class_val) class_number = tf.placeholder(tf.int32, shape=()) IoU_summary_output = tf.gather(IoU_summary_output_intermed, class_number) IoU_summary_no_reset_output = tf.gather(IoU_summary_no_reset_output_intermed, class_number) accuracy_per_class_output = tf.gather(accuracy_per_class_output_intermed, class_number) tf.summary.scalar("IoU per class", IoU_summary_output, collections=['per_class']) tf.summary.scalar("IoU (no reset) per class", IoU_summary_no_reset_output, collections=['per_class']) tf.summary.scalar("Accuracy per class", accuracy_per_class_output, collections=['per_class']) mIoU_output = tf.cond(mode, lambda: mIoU, lambda: Val_mIoU) mIoU_no_reset_output = tf.cond(mode, lambda: mIoU_no_reset, lambda: Val_mIoU_no_reset) tf.summary.scalar("mIoU", mIoU_output, collections=['all']) tf.summary.scalar("mIoU no reset", mIoU_no_reset_output, collections=['all']) tf.summary.image('images', tf.concat(axis=2, values=[images_summary, labels_summary, preds_summary]), max_outputs=args.save_num_images, collections=['all']) # Concatenate row-wise. summary_writer_train = tf.summary.FileWriter(os.path.join(args.snapshot_dir, 'train_all'), graph=tf.get_default_graph()) summary_writer_val = tf.summary.FileWriter(os.path.join(args.snapshot_dir, 'val_all'), graph=tf.get_default_graph()) summary_writer_per_class_val = [] summary_writer_per_class_train = [] for i in xrange(args.num_classes): summary_writer_per_class_train.append( tf.summary.FileWriter(os.path.join(args.snapshot_dir, 'train_class_' + str(i)), graph=tf.get_default_graph())) summary_writer_per_class_val.append( tf.summary.FileWriter(os.path.join(args.snapshot_dir, 'val_class_' + str(i)), graph=tf.get_default_graph())) # Define loss and optimisation parameters. base_lr = tf.constant(args.learning_rate) learning_rate = tf.scalar_mul(base_lr, tf.pow((1 - step_ph / args.num_steps), args.power)) tf.summary.scalar("learning_rate", learning_rate, collections=['all']) all_summary = tf.summary.merge_all('all') per_class_summary = tf.summary.merge_all('per_class') 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. sess = tf.Session() init = tf.group(tf.global_variables_initializer(), tf.local_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) # Load variables if the checkpoint is provided. if args.restore_from is not None: loader = tf.train.Saver(var_list=restore_var) load(loader, sess, args.restore_from) # Start queue threads. threads = tf.train.start_queue_runners(coord=coord, sess=sess) # Iterate over training steps. for step in xrange(1, args.num_steps + 1): start_time = time.time() # mode False -> val, mode True -> train if step % args.save_pred_every == 0: feed_dict = {step_ph: step, mode: False, class_number: step % args.num_classes} acc, loss_value, mI, mINR, _, _, _, summary_v_this_class, summary_v = sess.run( [accuracy_output, loss_output, mIoU_output, mIoU_no_reset_output, accuracy_per_class_output, IoU_summary_output, IoU_summary_no_reset_output, per_class_summary, all_summary], feed_dict=feed_dict) save(saver, sess, args.snapshot_dir, step) summary_writer_val.add_summary(summary_v, step) summary_writer_per_class_val[step % args.num_classes].add_summary(summary_v_this_class, step) duration = time.time() - start_time print( 'step {:d} \t Val_loss = {:.3f}, Val_acc = {:.3f}, Val_mIoU = {:.6f}, Val_mIoU_no_reset = {:.6f}, ({:.3f} sec/step)'.format( step, loss_value, acc, mI, mINR, duration)) else: feed_dict = {step_ph: step, mode: True, class_number: step % args.num_classes} acc, loss_value, mI, mINR, _, _, _, summary_t_this_class, summary_t, _ = sess.run( [accuracy_output, loss_output, mIoU_output, mIoU_no_reset_output, accuracy_per_class_output, IoU_summary_output, IoU_summary_no_reset_output, per_class_summary, all_summary, train_op], feed_dict=feed_dict) summary_writer_train.add_summary(summary_t, step) summary_writer_per_class_train[step % args.num_classes].add_summary(summary_t_this_class, step) duration = time.time() - start_time print( 'step {:d} \t loss = {:.3f}, acc = {:.3f}, mIoU = {:.6f}, mIoU_no_reset = {:.6f}, ({:.3f} sec/step)'.format( step, loss_value, acc, mI, mINR, duration)) coord.request_stop() # tboard_proc.kill() coord.join(threads)
def main(): """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) # Create queue coordinator. coord = tf.train.Coordinator() # Load reader. with tf.name_scope("create_inputs"): reader = ImageReader( args.data_dir, 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) image_batch075 = tf.image.resize_images(image_batch, [int(h * 0.75), int(w * 0.75)]) image_batch05 = tf.image.resize_images(image_batch, [int(h * 0.5), int(w * 0.5)]) # Create network. with tf.variable_scope('', reuse=False): net = DeepLabResNetModel({'data': image_batch}, is_training=args.is_training, num_classes=args.num_classes) with tf.variable_scope('', reuse=True): net075 = DeepLabResNetModel({'data': image_batch075}, is_training=args.is_training, num_classes=args.num_classes) with tf.variable_scope('', reuse=True): net05 = DeepLabResNetModel({'data': image_batch05}, is_training=args.is_training, num_classes=args.num_classes) # For a small batch size, it is better to keep # the statistics of the BN layers (running means and variances) # frozen, and to not update the values provided by the pre-trained model. # If is_training=True, the statistics will be updated during the training. # Note that is_training=False still updates BN parameters gamma (scale) and beta (offset) # if they are presented in var_list of the optimiser definition. # Predictions. raw_output100 = net.layers['fc1_voc12'] raw_output075 = net075.layers['fc1_voc12'] raw_output05 = net05.layers['fc1_voc12'] raw_output = tf.reduce_max(tf.stack([raw_output100, tf.image.resize_images(raw_output075, tf.shape(raw_output100)[1:3,]), tf.image.resize_images(raw_output05, tf.shape(raw_output100)[1:3,])]), axis=0) # Which variables to load. Running means and variances are not trainable, # thus all_variables() should be restored. restore_var = [v for v in tf.global_variables() if 'fc' not in v.name or not args.not_restore_last] all_trainable = [v for v in tf.trainable_variables() if 'beta' not in v.name and 'gamma' not in v.name] fc_trainable = [v for v in all_trainable if 'fc' in v.name] conv_trainable = [v for v in all_trainable if 'fc' not in v.name] # 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]) raw_prediction100 = tf.reshape(raw_output100, [-1, args.num_classes]) raw_prediction075 = tf.reshape(raw_output075, [-1, args.num_classes]) raw_prediction05 = tf.reshape(raw_output05, [-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] label_proc075 = prepare_label(label_batch, tf.stack(raw_output075.get_shape()[1:3]), num_classes=args.num_classes, one_hot=False) label_proc05 = prepare_label(label_batch, tf.stack(raw_output05.get_shape()[1:3]), num_classes=args.num_classes, one_hot=False) raw_gt = tf.reshape(label_proc, [-1,]) raw_gt075 = tf.reshape(label_proc075, [-1,]) raw_gt05 = tf.reshape(label_proc05, [-1,]) indices = tf.squeeze(tf.where(tf.less_equal(raw_gt, args.num_classes - 1)), 1) indices075 = tf.squeeze(tf.where(tf.less_equal(raw_gt075, args.num_classes - 1)), 1) indices05 = tf.squeeze(tf.where(tf.less_equal(raw_gt05, args.num_classes - 1)), 1) gt = tf.cast(tf.gather(raw_gt, indices), tf.int32) gt075 = tf.cast(tf.gather(raw_gt075, indices075), tf.int32) gt05 = tf.cast(tf.gather(raw_gt05, indices05), tf.int32) prediction = tf.gather(raw_prediction, indices) prediction100 = tf.gather(raw_prediction100, indices) prediction075 = tf.gather(raw_prediction075, indices075) prediction05 = tf.gather(raw_prediction05, indices05) # Pixel-wise softmax loss. loss = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=prediction, labels=gt) loss100 = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=prediction100, labels=gt) loss075 = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=prediction075, labels=gt075) loss05 = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=prediction05, labels=gt05) 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.reduce_mean(loss100) + tf.reduce_mean(loss075) + tf.reduce_mean(loss05) + 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()) # Define loss and optimisation parameters. 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)) 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) # Define a variable to accumulate gradients. accum_grads = [tf.Variable(tf.zeros_like(v.initialized_value()), trainable=False) for v in conv_trainable + fc_w_trainable + fc_b_trainable] # Define an operation to clear the accumulated gradients for next batch. zero_op = [v.assign(tf.zeros_like(v)) for v in accum_grads] # Compute gradients. grads = tf.gradients(reduced_loss, conv_trainable + fc_w_trainable + fc_b_trainable) # Accumulate and normalise the gradients. accum_grads_op = [accum_grads[i].assign_add(grad / args.grad_update_every) for i, grad in enumerate(grads)] grads_conv = accum_grads[:len(conv_trainable)] grads_fc_w = accum_grads[len(conv_trainable) : (len(conv_trainable) + len(fc_w_trainable))] grads_fc_b = accum_grads[(len(conv_trainable) + len(fc_w_trainable)):] # Apply the gradients. 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 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) # Load variables if the checkpoint is provided. if args.restore_from is not None: loader = tf.train.Saver(var_list=restore_var) load(loader, sess, args.restore_from) # 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 } loss_value = 0 # Clear the accumulated gradients. sess.run(zero_op, feed_dict=feed_dict) # Accumulate gradients. for i in range(args.grad_update_every): _, l_val = sess.run([accum_grads_op, reduced_loss], feed_dict=feed_dict) loss_value += l_val # Normalise the loss. loss_value /= args.grad_update_every # Apply gradients. if step % args.save_pred_every == 0: images, labels, summary, _ = sess.run([image_batch, label_batch, total_summary, train_op], feed_dict=feed_dict) summary_writer.add_summary(summary, step) save(saver, sess, args.snapshot_dir, step) else: sess.run(train_op, feed_dict=feed_dict) duration = time.time() - start_time print('step {:d} \t loss = {:.3f}, ({:.3f} sec/step)'.format(step, loss_value, duration)) coord.request_stop() coord.join(threads)
def main(): """Create the model and start the training.""" args = get_arguments() h, w = map(int, args.input_size.split(',')) input_size = (h, w) # Create queue coordinator. coord = tf.train.Coordinator() # Load reader. with tf.name_scope("create_inputs"): reader = ImageReader(args.data_dir, args.data_list, input_size, args.random_scale, coord) image_batch, label_batch = reader.dequeue(args.batch_size) # Create network. net = DeepLabResNetModel({'data': image_batch}, is_training=args.is_training) # For a small batch size, it is better to keep # the statistics of the BN layers (running means and variances) # frozen, and to not update the values provided by the pre-trained model. # If is_training=True, the statistics will be updated during the training. # Note that is_training=False still updates BN parameters gamma (scale) and beta (offset) # if they are presented in var_list of the optimiser definition. # Predictions. raw_output = net.layers['fc1_voc12'] # Which variables to load. Running means and variances are not trainable, # thus all_variables() should be restored. restore_var = tf.global_variables() trainable = tf.trainable_variables() prediction = tf.reshape(raw_output, [-1, n_classes]) label_proc = prepare_label(label_batch, tf.pack(raw_output.get_shape()[1:3])) gt = tf.reshape(label_proc, [-1, n_classes]) # Pixel-wise softmax loss. loss = tf.nn.softmax_cross_entropy_with_logits(prediction, gt) reduced_loss = tf.reduce_mean(loss) # Processed predictions. 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], tf.uint8) labels_summary = tf.py_func(decode_labels, [label_batch, args.save_num_images], tf.uint8) preds_summary = tf.py_func(decode_labels, [pred, args.save_num_images], tf.uint8) total_summary = tf.summary.image( 'images', tf.concat(2, [images_summary, labels_summary, preds_summary]), max_outputs=args.save_num_images) # Concatenate row-wise. summary_writer = tf.summary.FileWriter(args.snapshot_dir) # Define loss and optimisation parameters. optimiser = tf.train.AdamOptimizer(learning_rate=args.learning_rate) optim = optimiser.minimize(reduced_loss, var_list=trainable) # Set up tf session and initialize variables. config = tf.ConfigProto() config.gpu_options.allow_growth = True 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=restore_var, max_to_keep=40) # Load variables if the checkpoint is provided. if args.restore_from is not None: loader = tf.train.Saver(var_list=restore_var) load(loader, sess, args.restore_from) # 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() if step % args.save_pred_every == 0: loss_value, images, labels, preds, summary, _ = sess.run([ reduced_loss, image_batch, label_batch, pred, total_summary, optim ]) summary_writer.add_summary(summary, step) save(saver, sess, args.snapshot_dir, step) else: loss_value, _ = sess.run([reduced_loss, optim]) duration = time.time() - start_time print('step {:d} \t loss = {:.3f}, ({:.3f} sec/step)'.format( step, loss_value, duration)) coord.request_stop() coord.join(threads)
def main(): """Create the model and start the training.""" args = get_arguments() IMG_MEAN = np.zeros(3) image_std = [1.0, 1.0, 1.0] # parameters of building data set citylist = [ 'Norfolk', 'Arlington', 'Atlanta', 'Austin', 'Seekonk', 'NewHaven' ] image_mean_list = { 'Norfolk': [127.07435926, 129.40160709, 128.28713284], 'Arlington': [88.30304996, 94.97338776, 93.21268212], 'Atlanta': [101.997014375, 108.42171833, 110.044871], 'Austin': [97.0896012682, 102.94697026, 100.7540157], 'Seekonk': [86.67800904, 93.31221168, 92.1328146], 'NewHaven': [106.7092798, 111.4314, 110.74903832] } # BGR mean for the training data for each city image_std_list = { 'Norfolk': [28.615469420031832, 32.662536832452886, 37.64149854207523], 'Arlington': [30.40903039206398, 37.973725024862595, 43.58402191634698], 'Atlanta': [36.93662467838125, 39.43470059838385, 41.74732676809388], 'Austin': [42.81337177109884, 43.71071321350751, 44.440517007230675], 'Seekonk': [25.506449467410715, 32.46885262572024, 37.76814267502958], 'NewHaven': [33.05784541012469, 36.62685162291547, 37.686084270914435] } # set evaluation data if args.training_data == 'SP': IMG_MEAN = np.array( (121.68045527, 132.14961763, 129.30317439), dtype=np.float32) # mean of solar panel data in BGR order elif args.training_data in citylist: print("Training on {} data".format(args.training_data)) IMG_MEAN = image_mean_list[args.training_data] if args.unit_std: image_std = image_std_list[args.training_data] else: print("Wrong data option: {}".format(args.data_option)) # setup used GPU os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" os.environ["CUDA_VISIBLE_DEVICES"] = args.GPU h, w = map(int, args.input_size.split(',')) input_size = (h, w) tf.set_random_seed(args.random_seed) # Create queue coordinator. coord = tf.train.Coordinator() # Load reader. with tf.name_scope("create_inputs"): reader = ImageReader(args.data_dir, args.data_list, input_size, args.random_scale, args.random_mirror, args.random_rotate, args.ignore_label, IMG_MEAN, coord, image_std) image_batch, label_batch = reader.dequeue(args.batch_size) # Create network. net = DeepLabResNetModel({'data': image_batch}, is_training=args.is_training, num_classes=args.num_classes) # For a small batch size, it is better to keep # the statistics of the BN layers (running means and variances) # frozen, and to not update the values provided by the pre-trained model. # If is_training=True, the statistics will be updated during the training. # Note that is_training=False still updates BN parameters gamma (scale) and beta (offset) # if they are presented in var_list of the optimiser definition. # Predictions. raw_output = net.layers['fc1_voc12'] # Which variables to load. Running means and variances are not trainable, # thus all_variables() should be restored. restore_var = [ v for v in tf.global_variables() if 'fc' not in v.name or not args.not_restore_last ] all_trainable = [ v for v in tf.trainable_variables() if 'beta' not in v.name and 'gamma' not in v.name ] fc_trainable = [v for v in all_trainable if 'fc' in v.name] conv_trainable = [v for v in all_trainable if 'fc' not in v.name] # 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) # loss = tf.nn.softmax_cross_entropy_with_logits(logits=prediction, labels=gt) l2_losses = [ tf.nn.l2_loss(v) for v in tf.trainable_variables() if 'weights' in v.name ] reduced_loss = tf.reduce_mean( loss) + args.weight_decay * 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) # Define loss and optimisation parameters. base_lr = tf.constant(args.learning_rate) # step_ph = tf.placeholder(dtype=tf.float32, shape=()) global_step = tf.Variable(0, dtype=tf.int32, trainable=False, name='global_step') global_epoch = tf.cast(global_step * args.batch_size / args.data_size, tf.int8) # learning_rate = tf.scalar_mul(base_lr, tf.pow((1 - step_ph / args.num_steps), args.power)) if (args.decay_step <= 0): learning_rate = base_lr else: learning_rate = tf.train.exponential_decay( base_lr, global_step, tf.cast(args.data_size / args.batch_size * args.decay_step, tf.int32), args.decay_rate, staircase=True) 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), global_step=global_step) 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) # IOU of training batch # IOU = IOU_tf(pred, label_batch) # Image summary. images_summary = tf.py_func( inv_preprocess, [image_batch, args.save_num_images, IMG_MEAN, image_std], 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) image_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. loss_summary = tf.summary.scalar("loss", reduced_loss) entropy_summary = tf.summary.scalar("entropy", tf.reduce_mean(loss)) l2_loss_summary = tf.summary.scalar("L2_loss", tf.add_n(l2_losses)) learning_rate_summary = tf.summary.scalar( "learning_rate", learning_rate) # summary recording learning rate # IOU_summary = tf.summary.scalar("IOU", IOU) summary_writer = tf.summary.FileWriter(args.snapshot_dir, graph=tf.get_default_graph()) print("Setting up summary op...") total_summary = tf.summary.merge_all() # Set up tf session and initialize variables. config = tf.ConfigProto() config.gpu_options.allow_growth = True sess = tf.Session(config=config) init = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) sess.run(init) # Saver for storing checkpoints of the model. saver = tf.train.Saver(var_list=tf.global_variables(), max_to_keep=2) # saver for save/load checkpoint # load weights from saved checkpoint or initial pre-trained model if os.path.isdir(args.snapshot_dir): # search checkpoint at given path ckpt = tf.train.get_checkpoint_state(args.snapshot_dir) if ckpt and ckpt.model_checkpoint_path: # load checkpoint file load(saver, sess, ckpt.model_checkpoint_path) elif os.path.isfile(args.snapshot_dir): # load checkpoint file load(saver, sess, args.snapshot_dir) elif args.restore_from is not None: loader = tf.train.Saver( var_list=restore_var) # loader for part of pre-trained model load(loader, sess, args.restore_from) print("Load weights from{}".format(args.restore_from)) else: print("No model found at{}".format(args.restore_from)) # Start queue threads. threads = tf.train.start_queue_runners(coord=coord, sess=sess) # Initial status loss_value, entropy_loss, summary, itr, epoch = sess.run([ reduced_loss, tf.reduce_mean(loss), total_summary, global_step, global_epoch ]) print('step {:d} \t loss = {:.3f}, entropy_loss = {:.3f})'.format( itr, loss_value, entropy_loss)) summary_writer.add_summary(summary, itr) # Iterate over training steps. while (epoch < args.num_epochs): start_time = time.time() # feed_dict = { step_ph : step } _, itr, epoch = sess.run([train_op, global_step, global_epoch]) # save summary file if itr % 100 == 0: duration = time.time() - start_time pred_temp, truth_temp, loss_value, entropy_loss, summary, itr = sess.run( [ pred, label_batch, reduced_loss, tf.reduce_mean(loss), total_summary, global_step ]) summary_writer.add_summary(summary, itr) IOU_temp = IOU_(pred_temp, truth_temp) print( 'Epoch {:d} step {:d} \t loss = {:.3f}, entropy_loss = {:.3f}, IOU = {:.3f}, ({:.3f} sec/step)' .format(epoch, itr, loss_value, entropy_loss, IOU_temp, duration)) # print('Epoch {:d} step {:d} \t loss = {:.3f}, entropy_loss = {:.3f}, ({:.3f} sec/step)'.format( # epoch, itr, loss_value, entropy_loss, duration)) # save checkpoint if itr % args.save_pred_every == 0: # images, labels, preds = sess.run([image_batch, label_batch, pred]) save(saver, sess, args.snapshot_dir, global_step) # final status loss_value, entropy_loss, summary, itr = sess.run( [reduced_loss, tf.reduce_mean(loss), total_summary, global_step]) print('step {:d} \t loss = {:.3f}, entropy_loss = {:.3f}'.format( itr, loss_value, entropy_loss)) save(saver, sess, args.snapshot_dir, global_step) summary_writer.add_summary(summary, itr) coord.request_stop() coord.join(threads) sess.close()
def main(): """Create the model and start the evaluation process.""" args = get_arguments() h, w = map(int, args.input_size.split(',')) input_size = (h, w) # Create queue coordinator. coord = tf.train.Coordinator() # Load reader. with tf.name_scope("create_inputs"): reader = ImageReader( args.data_dir, args.data_list, input_size, # No defined input size. False, # No random scale. False, # No random mirror. args.ignore_label, IMG_MEAN, coord, ADV_FLAG, MASK_FLAG, args.eps, args.attack, args.targeted) # image, label = reader.image, reader.label image_batch, label_batch = reader.dequeue(args.batch_size) # image_batch, label_batch = tf.expand_dims(image, dim=0), tf.expand_dims(label, dim=0) # Add one batch dimension. image_name_list = reader.image_list print(image_name_list[0]) # Create network. net = DeepLabResNetModel({'data': image_batch}, is_training=False, num_classes=args.num_classes) # Which variables to load. restore_var = tf.global_variables() # Predictions. raw_output = net.layers['fc1_voc12'] 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_0 = tf.squeeze( tf.where(tf.less_equal(raw_gt, args.num_classes - 1)), 1) gt_0 = tf.cast(tf.gather(raw_gt, indices_0), tf.int32) prediction = tf.gather(raw_prediction, indices_0) raw_output = tf.image.resize_bilinear(raw_output, tf.shape(image_batch)[1:3, ]) raw_output = tf.argmax(raw_output, dimension=3) pred = tf.expand_dims(raw_output, dim=3) # Create 4-d tensor. # mIoU pred = tf.reshape(pred, [ -1, ]) gt = tf.reshape(label_batch, [ -1, ]) indices = tf.squeeze(tf.where(tf.less_equal(gt, args.num_classes - 1)), 1) ## ignore all labels >= num_classes gt = tf.cast(tf.gather(gt, indices), tf.int32) pred = tf.gather(pred, indices) mIoU, update_op = tf.contrib.metrics.streaming_mean_iou( pred, gt, num_classes=args.num_classes) # Set up tf session and initialize variables. config = tf.ConfigProto() config.gpu_options.allow_growth = True sess = tf.Session(config=config) init = tf.global_variables_initializer() sess.run(init) sess.run(tf.local_variables_initializer()) # Load weights. loader = tf.train.Saver(var_list=restore_var) if args.restore_from is not None: load(loader, sess, args.restore_from) # Start queue threads. threads = tf.train.start_queue_runners(coord=coord, sess=sess) # Iterate over evaluation steps. for step in range(args.num_steps): preds, _, X, Y = sess.run([pred, update_op, image_batch, label_batch]) if step % 100 == 0: print('step {:d}'.format(step)) print('Mean IoU: {:.3f}'.format(mIoU.eval(session=sess))) coord.request_stop() coord.join(threads)
def inference(self, i1_g, i2_g, d1_g, d2_g, f12_g, v1_g, m1_g, rt12_g, s1_g, off_h, off_w, fy, fx, y0, x0, is_train=True, reuse=False): with tf.variable_scope("inference"): if hyp.do_depth: # CHA adding valid depth_stuff = DepthNet(i1_g, d1_g, hyp.valid, is_train=(is_train and hyp.do_train_depth), reuse=reuse) else: depth_stuff = 0 if hyp.do_flow: flow_stuff = FlowNet(i1_g, i2_g, f12_g, v1_g, is_train=(is_train and hyp.do_train_flow), reuse=reuse) else: flow_stuff = 0 if hyp.do_odo: # CHA _, d1_e, _ = depth_stuff Z1_e = tf.exp(d1_e) # Z1_e = tf.stop_gradient(Z1_e) Z1_g = tf.exp(d1_g) if hyp.do_flow: _, f12_e, _ = flow_stuff f12_e = tf.stop_gradient( f12_e) / 50.0 # flow has a pretty high mag if hyp.cat_rgb: odo_inputs = tf.concat(3, [i1_g, i2_g]) if hyp.do_flow and hyp.cat_flow: odo_inputs = tf.concat(3, [odo_inputs, f12_e]) else: # we have to use flow odo_inputs = f12_e if hyp.cat_angles: ang_xy = tf.expand_dims( angleGrid(hyp.bs, hyp.h, hyp.w, y0, x0), 3) odo_inputs = tf.concat(3, [odo_inputs, ang_xy]) if hyp.do_flow and hyp.cat_ang_flow: x, y = f12_e[:, :, :, 0], f12_e[:, :, :, 1] ang_flow = tf.expand_dims(atan2_ocv(y, x), 3) odo_inputs = tf.concat(3, [odo_inputs, ang_flow]) if hyp.do_flow and hyp.cat_ang_diff: x, y = f12_e[:, :, :, 0], f12_e[:, :, :, 1] ang_flow = tf.expand_dims(atan2_ocv(y, x), 3) ang_xy = tf.expand_dims( angleGrid(hyp.bs, hyp.h, hyp.w, y0, x0), 3) ang_diff = ang_flow - ang_xy odo_inputs = tf.concat(3, [odo_inputs, ang_diff]) if hyp.cat_depth: odo_inputs = tf.concat(3, [odo_inputs, Z1_e]) odo_stuff = OdoNet(odo_inputs, i1_g, i2_g, Z1_e, Z1_g, rt12_g, m1_g, off_h, off_w, fy, fx, y0, x0, is_train=(is_train and hyp.do_train_odo), reuse=reuse) else: odo_stuff = 0 if hyp.do_seg: with tf.variable_scope("seg"): with tf.variable_scope("model"): input_size = (hyp.h, hyp.w) net = DeepLabResNetModel( {'data': i1_g}, is_training=hyp.do_train_seg, num_classes=hyp.num_seg_classes) # Predictions. raw_output = net.layers['fc1_voc12'] # Predictions: ignoring all predictions with labels greater or equal than n_classes raw_prediction = tf.reshape(raw_output, [-1, hyp.num_seg_classes]) label_proc = prepare_label( s1_g, tf.stack(raw_output.get_shape()[1:3]), num_classes=hyp.num_seg_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, hyp.num_seg_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) # Processed predictions: for visualisation. raw_output_up = tf.image.resize_bilinear( raw_output, tf.shape(i1_g)[1:3, ]) raw_output_up_lab = tf.argmax(raw_output_up, dimension=3) pred = tf.expand_dims(raw_output_up_lab, dim=3) # deining the l2 loss on the parameters of the model weight_var = [ v for v in tf.trainable_variables() if 'inference/seg/model' in v.name ] weight_var = [v for v in weight_var if 'weights' in v.name] # l2_losses = [hyp.seg_weight_decay * tf.nn.l2_loss(v) for v in tf.trainable_variables() if 'weights' in v.name] l2_losses = [ hyp.seg_weight_decay * tf.nn.l2_loss(v) for v in weight_var ] reduced_loss = tf.reduce_mean(loss) + tf.add_n(l2_losses) seg_stuff = (reduced_loss, raw_output_up, pred) else: seg_stuff = 0 return depth_stuff, flow_stuff, odo_stuff, seg_stuff
def main(): os.environ['CUDA_VISIBLE_DEVICES'] = '3' """Create the model and start the training.""" args = get_arguments() # Get width (w) and height (h) of an input image h, w = map(int, args.input_size.split(',')) input_size = (h, w) # Set random seed for reproducibility of the results tf.set_random_seed(args.random_seed) # Create queue coordinator. coord = tf.train.Coordinator() # Load reader. with tf.name_scope("create_inputs"): reader = ImageReader( args.data_dir, 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) # 1st = images, 2nd = gt labels, 3rd = imu (we do not need IMU here) # Create network. with tf.variable_scope('', reuse=False): net = wasr_NOIMU2({'data': image_batch}, is_training=args.is_training, num_classes=args.num_classes) # For a small batch size, it is better to keep # the statistics of the BN layers (running means and variances) # frozen, and to not update the values provided by the pre-trained model. # If is_training=True, the statistics will be updated during the training. # Note that is_training=False still updates BN parameters gamma (scale) and beta (offset) # if they are presented in var_list of the optimiser definition. # Predictions. # The layer at which the final output is located raw_output = net.layers['fc1_voc12'] # The layer from which we extract features for computing water-obstacle separation loss inthemiddle_output = net.layers['res4b20'] # Which variables to load. Running means and variances are not trainable, # thus all_variables() should be restored. restore_var = [v for v in tf.global_variables() if 'fc' not in v.name or not args.not_restore_last] #all_trainable = [v for v in tf.trainable_variables() if 'beta' not in v.name and 'gamma'] # all trainable variables all_trainable = [v for v in tf.trainable_variables() if 'beta' not in v.name and 'gamma' not in v.name] #new all trainable fc_trainable = [v for v in all_trainable if 'fc' in v.name] # only variables from the ASPP module arm_trainable = [v for v in all_trainable if 'arm_conv' in v.name] # only variables from the ARM module ffm_trainable = [v for v in all_trainable if 'ffm_conv' in v.name] # only variables from the FFM module batch_trainable = [v for v in tf.trainable_variables() if 'beta' in v.name or 'gamma' in v.name] # for batchnorms conv_trainable = [v for v in all_trainable if 'fc' not in v.name and 'arm_conv' not in v.name and 'ffm_conv' not in v.name] # 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 # Check if everything sums up correctly. Do the neccessary assertions print("----") print("Number of all trainable: {:d}\nNumber of fc trainable: {:d}\nNumber of conv trainable: {:d}\nNumber of ARM trainable: {:d}\nNumber of FFM trainable: {:d}\n".format(len(all_trainable), len(fc_trainable), len(conv_trainable), len(arm_trainable), len(ffm_trainable))) assert(len(all_trainable) == len(fc_trainable) + len(conv_trainable) + len(arm_trainable) + len(ffm_trainable)) print("----") print("Number of fc trainable: {:d}\nNumber of fc_w trainable: {:d}\nNumber of fc_b trainable: {:d}\n".format(len(fc_trainable), len(fc_w_trainable), len(fc_b_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) 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) # Features loss from somewhere in the middle. This forces the network to separate water pixels from obstacles loss_0 = cost_function_separate_water_obstacle(inthemiddle_output, label_batch) #loss_0 = tf.Print(loss_0, [loss_0], 'Water separation loss ') # Pixel-wise softmax cross entropy loss (This is the tensorflow implementation) ce_loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=prediction, labels=gt)) ce_loss = tf.Print(ce_loss, [ce_loss], 'Default TF crossentropy loss ') # Weight decay losses (l2 regularization) l2_losses = [args.weight_decay * tf.nn.l2_loss(v) for v in tf.trainable_variables() if 'weights' in v.name] added_l2_losses = 10.e-2 * tf.add_n(l2_losses) # add together all l2 losses added_l2_losses = tf.Print(added_l2_losses, [added_l2_losses], message="l2 losses ") # Focal loss focal_loss = focal_loss_cost(labels=gt, logits=prediction) focal_loss = tf.Print(focal_loss, [focal_loss], message="Focal loss ") # Add together all of the losses (focal loss, weight decay and water-separation loss) reduced_loss = added_l2_losses + focal_loss + loss_0 #(10.e-2 * loss_0) # focal loss #reduced_loss = added_l2_losses + ce_loss + loss_0 #(10.e-2 * loss_0) # normal cross entropy # Define loss and optimisation parameters. base_lr = tf.constant(args.learning_rate) step_ph = tf.placeholder(dtype=tf.float32, shape=()) # Learning rate modified based on the the current step learning_rate = tf.scalar_mul(base_lr, tf.pow((1 - step_ph / args.num_steps), args.power)) # version 1 #learning_rate = tf.train.exponential_decay(base_lr, step_ph, 750, 0.7, staircase=True) # version 2 # RMSProp optimizer opt_conv = tf.train.RMSPropOptimizer(learning_rate=learning_rate, decay=0.9, momentum=args.momentum, centered=True, name='RMSProp_conv') opt_sp_w = tf.train.RMSPropOptimizer(learning_rate=learning_rate * 10, decay=0.9, momentum=args.momentum, centered=True, name='RMSProp_special_w') opt_sp_b = tf.train.RMSPropOptimizer(learning_rate=learning_rate * 20, decay=0.9, momentum=args.momentum, centered=True, name='RMSProp_special_b') # Momentum optimizer (original) #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) # Minimization of optimizers for specific trainable variables... op_c_all = opt_conv.minimize(reduced_loss, var_list=[conv_trainable, batch_trainable]) op_spc_w = opt_sp_w.minimize(reduced_loss, var_list=[fc_w_trainable, arm_trainable, ffm_trainable]) op_spc_b = opt_sp_b.minimize(reduced_loss, var_list=[fc_b_trainable]) train_op = tf.group(op_c_all, op_spc_w, op_spc_b) # Set up tf session and initialize variables. config = tf.ConfigProto() config.gpu_options.allow_growth = True 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=1) # Load variables if the checkpoint is provided. #if args.restore_from is not None: # loader = tf.train.Saver(var_list=restore_var) # load(loader, sess, args.restore_from) # RESTORE PARTIAL WEIGHTS (which are available) restored_vars = get_tensors_in_checkpoint_file(file_name=args.restore_from, restore_last_bool=args.not_restore_last) tensors_to_load = build_tensors_in_checkpoint_file(restored_vars) loader = tf.train.Saver(var_list=tensors_to_load) loader.restore(sess, args.restore_from) # 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 } loss_value, _ = sess.run([reduced_loss, train_op], feed_dict=feed_dict) if step % args.save_pred_every == 0: save(saver, sess, args.snapshot_dir, step) duration = time.time() - start_time print('step {:d} \t loss = {:.3f}, ({:.3f} sec/step)'.format(step, loss_value, duration)) # join threads coord.request_stop() coord.join(threads)
def main(): with tf.device('/cpu:0'): args = get_arguments() with tf.name_scope('create_inputs'): image_batch = tf.placeholder(dtype=tf.float32, shape=[args.batch_size, 256, 256, 3]) label_batch = tf.placeholder(dtype=tf.uint8, shape=[args.batch_size, 256, 256, 1]) with tf.device('/gpu:0'): #create network net = DeepLabResNetModel({'data': image_batch}, is_training=False, num_classes=args.num_classes) raw_output = net.layers['fc1_voc12'] restore_var = tf.global_variables() print('raw_outout:', raw_output) prediction = tf.reshape(raw_output, [-1, args.num_classes]) print('prediction:', prediction) with tf.device('/cpu:0'): label_proc = prepare_label(label_batch, tf.stack(raw_output.get_shape()[1:3]), num_classes=args.num_classes) print('label_proc:', label_batch) gt = tf.reshape(label_proc, [-1, args.num_classes]) print('gt:', gt) loss = tf.nn.softmax_cross_entropy_with_logits(logits=prediction, labels=gt) reduced_loss = tf.reduce_mean(loss) #Processed predictions raw_output_up = tf.image.resize_bilinear(raw_output, tf.shape(image_batch)[1:3, ]) raw_output_upp = tf.argmax(raw_output_up, dimension=3) pred = tf.expand_dims(raw_output_upp, dim=3) with tf.device('/cpu:0'): input_batch = tf.squeeze(label_batch, squeeze_dims=[3]) label_one_hot = tf.one_hot(input_batch, depth=args.num_classes) with tf.device('/cpu:0'): iou_loss = tl.cost.iou_coe(raw_output_up, label_one_hot, axis=[1, 2, 3]) # Set up tf session and initialize variables. config = tf.ConfigProto() config.gpu_options.allow_growth = True sess = tf.Session(config=config) init = tf.global_variables_initializer() sess.run(init) with tf.device('/cpu:0'): # Load variables if the checkpoint is provided. if args.restore_from is not None: loader = tf.train.Saver(var_list=restore_var) load(loader, sess, args.restore_from) x_train, y_train, shapes = get_one_minibatch(args.data_list, args.batch_size, data_augument, True) feed_dict = {image_batch: x_train, label_batch: y_train} loss_value, IoU, pred_maskers = sess.run( [reduced_loss, iou_loss, pred], feed_dict=feed_dict) print('loss:{}, IoU: {}\n'.format(loss_value, IoU)) print('pred_masker shape: ', pred_maskers.shape) for num, masker in enumerate(pred_maskers): pred_masker = np.squeeze(masker) orignal_masker = np.squeeze(y_train[num]) misc.imsave(r'D:\tmp\{}_orignal_masker.jpg'.format(num), orignal_masker * 255) misc.imsave(r'D:\tmp\{}_pred_masker.jpg'.format(num), pred_masker * 500)
def main(): """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) # Create queue coordinator. coord = tf.train.Coordinator() # Load reader. with tf.name_scope("create_inputs"): reader = ImageReader( args.data_dir, 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) # Create network. net = DeepLabResNetModel({'data': image_batch}, is_training=args.is_training, num_classes=args.num_classes) # For a small batch size, it is better to keep # the statistics of the BN layers (running means and variances) # frozen, and to not update the values provided by the pre-trained model. # If is_training=True, the statistics will be updated during the training. # Note that is_training=False still updates BN parameters gamma (scale) and beta (offset) # if they are presented in var_list of the optimiser definition. # Predictions. raw_output = net.layers['fc1_voc12'] # Which variables to load. Running means and variances are not trainable, # thus all_variables() should be restored. # Restore all variables, or all except the last ones. restore_var = [v for v in tf.global_variables() if 'fc' not in v.name or not args.not_restore_last] trainable = [v for v in tf.trainable_variables() if 'fc1_voc12' in v.name] # Fine-tune only the last layers. 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. if not args.class_weights: loss = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=prediction, labels=gt) # Multiply logits by appropriate class weight else: raw_weights = tf.gather(args.class_weights, tf.cast(raw_gt, tf.int32)) weights = tf.gather(raw_weights, indices) loss = tf.losses.sparse_softmax_cross_entropy(logits=prediction, labels=gt, weights=weights) reduced_loss = tf.reduce_mean(loss) # Processed predictions. 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()) # Define loss and optimisation parameters. optimiser = tf.train.AdamOptimizer(learning_rate=args.learning_rate) optim = optimiser.minimize(reduced_loss, var_list=trainable) # Prep val data if args.val_list: val_steps = int(args.val_size / args.batch_size) with tf.name_scope("get_val"): reader_val = ImageReader( args.data_dir, args.val_list, input_size, False, False, args.ignore_label, IMG_MEAN, coord) val_image_batch, val_label_batch = reader.dequeue(args.batch_size) # Val predictions. val_raw_output = tf.image.resize_bilinear(raw_output, tf.shape(val_image_batch)[1:3,]) val_raw_output = tf.argmax(val_raw_output, dimension=3) val_pred = tf.expand_dims(val_raw_output, dim=3) # Create 4-d tensor. # mIoU val_pred = tf.reshape(val_pred, [-1,]) val_gt = tf.reshape(val_label_batch, [-1,]) weights = tf.cast(tf.less_equal(val_gt, args.num_classes - 1), tf.int32) # Ignoring all labels greater than or equal to n_classes. mIoU, update_op = tf.contrib.metrics.streaming_mean_iou(val_pred, val_gt, num_classes=args.num_classes, weights=weights) # Set up tf session and initialize variables. config = tf.ConfigProto() config.gpu_options.allow_growth = True sess = tf.Session(config=config) init = tf.global_variables_initializer() sess.run(init) if args.val_list: sess.run(tf.local_variables_initializer()) # Saver for storing checkpoints of the model. saver = tf.train.Saver(var_list=tf.global_variables(), max_to_keep=20) # Load variables if the checkpoint is provided. if args.restore_from is not None: loader = tf.train.Saver(var_list=restore_var) load(loader, sess, args.restore_from) # 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() if step % args.save_pred_every == 0: loss_value, images, labels, preds, summary, _ = sess.run([reduced_loss, image_batch, label_batch, pred, total_summary, optim]) summary_writer.add_summary(summary, step) # Print val jaccard loss if args.val_list: for vstep in range(val_steps): val_preds, _ = sess.run([val_pred, update_op]) viou = mIoU.eval(session=sess) print('Mean IoU: {:.6f}'.format(viou)) save(saver, sess, args.snapshot_dir, step, val_iou=viou) else: save(saver, sess, args.snapshot_dir, step) else: loss_value, _ = sess.run([reduced_loss, optim]) duration = time.time() - start_time print('step {:d} \t loss = {:.3f}, ({:.3f} sec/step)'.format(step, loss_value, duration)) coord.request_stop() coord.join(threads)
def main(): """Create the model and start the training.""" args = get_arguments() h, w = map(int, args.input_size.split(',')) input_size = (h, w) # Create queue coordinator. coord = tf.train.Coordinator() # Load reader. with tf.name_scope("create_inputs"): reader = ImageReader(args.data_dir, args.data_list, input_size, RANDOM_SCALE, coord) image_batch, label_batch = reader.dequeue(args.batch_size) # Create network. net = DeepLabResNetModel({'data': image_batch}) # Predictions. raw_output = net.layers['fc1_voc12'] prediction = tf.reshape(raw_output, [-1, n_classes]) label_proc = prepare_label(label_batch, tf.pack(raw_output.get_shape()[1:3])) gt = tf.reshape(label_proc, [-1, n_classes]) # Pixel-wise softmax loss. loss = tf.nn.softmax_cross_entropy_with_logits(prediction, gt) reduced_loss = tf.reduce_mean(loss) # Processed predictions. 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) # Define loss and optimisation parameters. optimiser = tf.train.AdamOptimizer(learning_rate=args.learning_rate) trainable = tf.trainable_variables() optim = optimiser.minimize(reduced_loss, var_list=trainable) # Set up tf session and initialize variables. config = tf.ConfigProto() config.gpu_options.allow_growth = True sess = tf.Session(config=config) init = tf.initialize_all_variables() sess.run(init) # Saver for storing checkpoints of the model. saver = tf.train.Saver(var_list=trainable, max_to_keep=40) if args.restore_from is not None: load(saver, sess, args.restore_from) # Start queue threads. threads = tf.train.start_queue_runners(coord=coord, sess=sess) if not os.path.exists(args.save_dir): os.makedirs(args.save_dir) # Iterate over training steps. for step in range(args.num_steps): start_time = time.time() if step % args.save_pred_every == 0: loss_value, images, labels, preds, _ = sess.run( [reduced_loss, image_batch, label_batch, pred, optim]) fig, axes = plt.subplots(args.save_num_images, 3, figsize=(16, 12)) for i in xrange(args.save_num_images): axes.flat[i * 3].set_title('data') axes.flat[i * 3].imshow( (images[i] + IMG_MEAN)[:, :, ::-1].astype(np.uint8)) axes.flat[i * 3 + 1].set_title('mask') axes.flat[i * 3 + 1].imshow(decode_labels(labels[i, :, :, 0])) axes.flat[i * 3 + 2].set_title('pred') axes.flat[i * 3 + 2].imshow(decode_labels(preds[i, :, :, 0])) plt.savefig(args.save_dir + str(start_time) + ".png") plt.close(fig) save(saver, sess, args.snapshot_dir, step) else: loss_value, _ = sess.run([reduced_loss, optim]) duration = time.time() - start_time print('step {:d} \t loss = {:.3f}, ({:.3f} sec/step)'.format( step, loss_value, duration)) coord.request_stop() coord.join(threads)
def main(): """Create the model and start the training.""" args = get_arguments() # setup used GPU os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" os.environ["CUDA_VISIBLE_DEVICES"] = args.GPU h, w = map(int, args.input_size.split(',')) input_size = (h, w) tf.set_random_seed(args.random_seed) # Create queue coordinator. coord = tf.train.Coordinator() # Load reader. with tf.name_scope("create_inputs"): reader = ImageReader(args.data_dir, 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) # Create network. net = DeepLabResNetModel({'data': image_batch}, is_training=args.is_training, num_classes=args.num_classes) # For a small batch size, it is better to keep # the statistics of the BN layers (running means and variances) # frozen, and to not update the values provided by the pre-trained model. # If is_training=True, the statistics will be updated during the training. # Note that is_training=False still updates BN parameters gamma (scale) and beta (offset) # if they are presented in var_list of the optimiser definition. # Predictions. raw_output = net.layers['fc1_voc12'] # Which variables to load. Running means and variances are not trainable, # thus all_variables() should be restored. # Restore all variables, or all except the last ones. restore_var = [ v for v in tf.global_variables() if 'fc' not in v.name or not args.not_restore_last ] trainable = [v for v in tf.trainable_variables() if 'fc1_voc12' in v.name] # Fine-tune only the last layers. 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) gt = tf.reshape(label_proc, [-1, args.num_classes]) # Pixel-wise softmax loss. loss = tf.nn.softmax_cross_entropy_with_logits(logits=prediction, labels=gt) reduced_loss = tf.reduce_mean(loss) # Processed predictions. 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) # Define loss and optimisation parameters. global_step = tf.Variable(0, dtype=tf.int32, trainable=False, name='global_step') optimiser = tf.train.AdamOptimizer(learning_rate=args.learning_rate) optim = optimiser.minimize(reduced_loss, var_list=trainable, global_step=global_step) # 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) image_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. loss_summary = tf.summary.scalar("loss", reduced_loss) summary_writer = tf.summary.FileWriter(args.snapshot_dir, graph=tf.get_default_graph()) print("Setting up summary op...") total_summary = tf.summary.merge_all() # Set up tf session and initialize variables. config = tf.ConfigProto() config.gpu_options.allow_growth = True 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=40) # Load variables if the checkpoint is provided. ckpt = tf.train.get_checkpoint_state(args.snapshot_dir) if ckpt and ckpt.model_checkpoint_path: saver.restore(sess, ckpt.model_checkpoint_path) elif args.restore_from is not None: loader = tf.train.Saver(var_list=restore_var) load(loader, sess, args.restore_from) # Start queue threads. threads = tf.train.start_queue_runners(coord=coord, sess=sess) # Initial status loss_value, summary, itr = sess.run( [reduced_loss, total_summary, global_step]) print('step {:d} \t loss = {:.3f}'.format(itr, loss_value)) summary_writer.add_summary(summary, itr) # Iterate over training steps. for step in range(args.num_steps): start_time = time.time() _, itr = sess.run([optim, global_step]) # save summary file if itr % 100 == 0: duration = time.time() - start_time loss_value, summary = sess.run([reduced_loss, total_summary]) summary_writer.add_summary(summary, itr) print('step {:d} \t loss = {:.3f} ({:.3f} sec/step)'.format( itr, loss_value, duration)) # save checkpoint if itr % args.save_pred_every == 0: save(saver, sess, args.snapshot_dir, itr) # final status loss_value, summary, itr = sess.run( [reduced_loss, total_summary, global_step]) print('step {:d} \t loss = {:.3f}'.format(itr, loss_value)) save(saver, sess, args.snapshot_dir, global_step) summary_writer.add_summary(summary, itr) coord.request_stop() coord.join(threads)
def main(): """主函数:模型构建和训练""" # 获取训练参数 args = get_arguments() # 获取输出尺寸 h, w = map(int, args.input_size.split(',')) input_size = (h, w) # 设置随机种子 tf.set_random_seed(args.random_seed) # 构建队列协调器queue coordinator coord = tf.train.Coordinator() # 加载读取器reader. with tf.name_scope("create_inputs"): reader = ImageReader( args.data_dir, 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) # 构建DeepLab-ResNet-101网络 net = DeepLabResNetModel({'data': image_batch}, is_training=args.is_training, num_classes=args.num_classes) # 获取网络输出. raw_output = net.layers['fc1_voc12'] # 获取不同类型的网络权重参数名 restore_var = [v for v in tf.global_variables() if 'fc' not in v.name or not args.not_restore_last] all_trainable = [v for v in tf.trainable_variables() if 'beta' not in v.name and 'gamma' not in v.name] fc_trainable = [v for v in all_trainable if 'fc' in v.name] conv_trainable = [v for v in all_trainable if 'fc' not in v.name] # 学习率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)) # 根据ground truth类别标签,提取 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) # 构建各个像素的损失函数:交叉熵softmax loss + 权重衰减L2 loss # 交叉熵softmax loss loss = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=prediction, labels=gt) # 权重衰减L2 loss l2_losses = [args.weight_decay * tf.nn.l2_loss(v) for v in tf.trainable_variables() if 'weights' in v.name] # loss合并 reduced_loss = tf.reduce_mean(loss) + tf.add_n(l2_losses) # 预测结果可视化 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) # 添加图片总结 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()) # 定义loss和优化参数 # 定义学习率 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)) # 定义优化器 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) # 获取loss梯度 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) # 建立tf session config = tf.ConfigProto() config.gpu_options.allow_growth = True sess = tf.Session(config=config) # 执行权重变量初始化 init = tf.global_variables_initializer() sess.run(init) # 获取训练存储器 saver = tf.train.Saver(var_list=tf.global_variables(), max_to_keep=10) # 加载已有的checkpoint文件 if args.restore_from is not None: loader = tf.train.Saver(var_list=restore_var) load(loader, sess, args.restore_from) # 开启队列执行器线程. threads = tf.train.start_queue_runners(coord=coord, sess=sess) # 遍历所有训练step. 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, images, labels, preds, summary, _ = sess.run([reduced_loss, image_batch, label_batch, pred, total_summary, train_op], feed_dict=feed_dict) summary_writer.add_summary(summary, step) save(saver, sess, args.snapshot_dir, step) # 仅执行训练,不存储训练结果 else: loss_value, _ = sess.run([reduced_loss, train_op], feed_dict=feed_dict) duration = time.time() - start_time print('step {:d} \t loss = {:.3f}, ({:.3f} sec/step)'.format(step, loss_value, duration)) # 停止训练协调器 coord.request_stop() coord.join(threads)
def main(): """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) # Create queue coordinator. coord = tf.train.Coordinator() # Load reader. with tf.name_scope("create_inputs"): reader = ImageReader(args.data_dir, 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) # Create network. #net = DeepLabResNetModel({'data': image_batch}, is_training=args.is_training, num_classes=args.num_classes) net = DeepLabResNetModel({'data': image_batch}, is_training=False, num_classes=args.num_classes) # For a small batch size, it is better to keep # the statistics of the BN layers (running means and variances) # frozen, and to not update the values provided by the pre-trained model. # If is_training=True, the statistics will be updated during the training. # Note that is_training=False still updates BN parameters gamma (scale) and beta (offset) # if they are presented in var_list of the optimiser definition. # Predictions. raw_output = net.layers['fc1_voc12'] # Which variables to load. Running means and variances are not trainable, # thus all_variables() should be restored. # Restore all variables, or all except the last ones. 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) gt = tf.reshape(label_proc, [-1, args.num_classes]) # Pixel-wise softmax loss. loss = tf.nn.softmax_cross_entropy_with_logits(logits=prediction, labels=gt) reduced_loss = tf.reduce_mean(loss) # Processed predictions. 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) # Calculate IoU mean_IoU = tf.metrics.mean_iou(labels=label_batch, predictions=pred, num_classes=args.num_classes) mean_per_class_acc = tf.metrics.mean_per_class_accuracy( labels=label_batch, predictions=pred, num_classes=args.num_classes) mean_acc = tf.metrics.accuracy(labels=label_batch, predictions=pred) # 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) final_output = tf.concat( axis=2, values=[images_summary, labels_summary, preds_summary]) #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. #loss_summary = tf.summary.scalar('training_loss', reduced_loss) #overall_summary = tf.summary.merge([total_summary, loss_summary]) #summary_writer = tf.summary.FileWriter(args.snapshot_dir, # graph=tf.get_default_graph()) # Define loss and optimisation parameters. #optimiser = tf.train.AdamOptimizer(learning_rate=args.learning_rate) #optim = optimiser.minimize(reduced_loss, var_list=trainable) # Set up tf session and initialize variables. config = tf.ConfigProto() config.gpu_options.allow_growth = True sess = tf.Session(config=config) init = tf.global_variables_initializer() sess.run(init) sess.run(tf.local_variables_initializer()) # Saver for storing checkpoints of the model. saver = tf.train.Saver(var_list=tf.global_variables(), max_to_keep=40) saver.restore(sess, args.restore_from) # Load variables if the checkpoint is provided. #if args.restore_from is not None: # loader = tf.train.Saver(var_list=restore_var) # load(loader, sess, args.restore_from) # Start queue threads. threads = tf.train.start_queue_runners(coord=coord, sess=sess) start_time = time.time() for step in range(25): loss, IoU, per_class_acc, acc, output = sess.run([ reduced_loss, mean_IoU, mean_per_class_acc, mean_acc, final_output ]) ipdb.set_trace() print(IoU, per_class_acc, acc) ''' # Iterate over training steps. for step in range(args.num_steps): start_time = time.time() loss, image, overall_output = sess.run([reduced_loss, image_batch, final_output]) #if step % args.save_pred_every == 0: # loss_value, images, labels, preds, summary, _ = sess.run([reduced_loss, image_batch, label_batch, pred, overall_summary, optim]) # summary_writer.add_summary(summary, step) # save(saver, sess, args.snapshot_dir, step) #else: # loss_value, _, summary = sess.run([reduced_loss, optim, loss_summary]) # summary_writer.add_summary(summary, step) duration = time.time() - start_time #print('step {:d} \t loss = {:.3f}, ({:.3f} sec/step)'.format(step, loss_value, duration)) ''' coord.request_stop() coord.join(threads)
def main(): """Create the model and start the evaluation process.""" args = get_arguments() h, w = map(int, args.input_size.split(',')) input_size = (h, w) # Create queue coordinator. coord = tf.train.Coordinator() # Load reader. with tf.name_scope("create_inputs"): reader = ImageReader( args.data_dir, args.data_list, input_size, # No defined input size. False, # No random scale. False, # No random mirror. args.ignore_label, IMG_MEAN, coord) # image, label = reader.image, reader.label image_batch, label_batch = reader.dequeue(args.batch_size) # image_batch, label_batch = tf.expand_dims(image, dim=0), tf.expand_dims(label, dim=0) # Add one batch dimension. image_name_list = reader.image_list x = tf.placeholder(tf.float32, shape=(BATCH_SIZE, h, w, 3)) y = tf.placeholder(tf.float32, shape=(BATCH_SIZE, h, w, 1)) # Create network. net = DeepLabResNetModel({'data': x}, is_training=False, num_classes=args.num_classes) # Which variables to load. restore_var = tf.global_variables() # Predictions. raw_output = net.layers['fc1_voc12'] raw_prediction = tf.reshape(raw_output, [-1, args.num_classes]) label_proc = prepare_label(y, 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_0 = tf.squeeze( tf.where(tf.less_equal(raw_gt, args.num_classes - 1)), 1) gt_0 = tf.cast(tf.gather(raw_gt, indices_0), tf.int32) prediction = tf.gather(raw_prediction, indices_0) loss = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=prediction, labels=gt_0) img_mean_t = tf.convert_to_tensor(IMG_MEAN, dtype=tf.float32) if args.attack == 'fgs': x_adv = fgs(x, loss, args.eps, img_mean_t, input_size, BATCH_SIZE, MASK_FLAG, args.targeted) elif args.attack == 'ifgs': x_adv = fgs(x, loss, args.beta, img_mean_t, input_size, BATCH_SIZE, MASK_FLAG, args.targeted) raw_output = tf.image.resize_bilinear(raw_output, tf.shape(image_batch)[1:3, ]) raw_output = tf.argmax(raw_output, dimension=3) pred = tf.expand_dims(raw_output, dim=3) # Create 4-d tensor. # mIoU pred = tf.reshape(pred, [ -1, ]) gt = tf.reshape(label_batch, [ -1, ]) indices = tf.squeeze(tf.where(tf.less_equal(gt, args.num_classes - 1)), 1) ## ignore all labels >= num_classes gt = tf.cast(tf.gather(gt, indices), tf.int32) pred = tf.gather(pred, indices) mIoU, update_op = tf.contrib.metrics.streaming_mean_iou( pred, gt, num_classes=args.num_classes) # Set up tf session and initialize variables. config = tf.ConfigProto() config.gpu_options.allow_growth = True sess = tf.Session(config=config) init = tf.global_variables_initializer() sess.run(init) sess.run(tf.local_variables_initializer()) # Load weights. loader = tf.train.Saver(var_list=restore_var) if args.restore_from is not None: load(loader, sess, args.restore_from) # Start queue threads. threads = tf.train.start_queue_runners(coord=coord, sess=sess) # Iterate over evaluation steps. image_list, label_list = read_labeled_image_list(args.data_dir, args.data_list) for step in range(args.num_steps): X = image_converter(image_list, step) if args.targeted: Y = np.zeros((BATCH_SIZE, 321, 321, 1)) else: Y = label_converter(label_list, step) if args.attack == 'fgs': preds, _, X_adv = sess.run([pred, update_op, x_adv], feed_dict={ x: X, y: Y }) elif args.attack == 'ifgs': X_adv = X for i in range(args.iter): preds, _, X_adv, loss_v = sess.run( [pred, update_op, x_adv, loss], feed_dict={ x: X_adv, y: Y }) r = np.clip(X_adv - X, -args.eps, args.eps) X_adv = X + r # preds, _ = sess.run([pred, update_op], feed_dict={x: X_adv}) if SAVE_FLAG is not None: image_saver(X_adv, X, image_name_list[step], args.attack, args.eps, args.targeted) if step % 100 == 0: print('step {:d}'.format(step)) print('Mean IoU: {:.3f}'.format(mIoU.eval(session=sess))) coord.request_stop() coord.join(threads)
def main(): """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) # Create queue coordinator. coord = tf.train.Coordinator() # Load reader. with tf.name_scope("create_inputs"): reader = ImageReader(args.data_dir, 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) image_batch075 = tf.image.resize_images( image_batch, [int(h * 0.75), int(w * 0.75)]) image_batch05 = tf.image.resize_images( image_batch, [int(h * 0.5), int(w * 0.5)]) # Create network. with tf.variable_scope('', reuse=False): net = DeepLabResNetModel({'data': image_batch}, is_training=args.is_training, num_classes=args.num_classes) with tf.variable_scope('', reuse=True): net075 = DeepLabResNetModel({'data': image_batch075}, is_training=args.is_training, num_classes=args.num_classes) with tf.variable_scope('', reuse=True): net05 = DeepLabResNetModel({'data': image_batch05}, is_training=args.is_training, num_classes=args.num_classes) # For a small batch size, it is better to keep # the statistics of the BN layers (running means and variances) # frozen, and to not update the values provided by the pre-trained model. # If is_training=True, the statistics will be updated during the training. # Note that is_training=False still updates BN parameters gamma (scale) and beta (offset) # if they are presented in var_list of the optimiser definition. # Predictions. raw_output100 = net.layers['fc1_voc12'] raw_output075 = net075.layers['fc1_voc12'] raw_output05 = net05.layers['fc1_voc12'] raw_output = tf.reduce_max(tf.stack([ raw_output100, tf.image.resize_images(raw_output075, tf.shape(raw_output100)[1:3, ]), tf.image.resize_images(raw_output05, tf.shape(raw_output100)[1:3, ]) ]), axis=0) # Which variables to load. Running means and variances are not trainable, # thus all_variables() should be restored. restore_var = [ v for v in tf.global_variables() if 'fc' not in v.name or not args.not_restore_last ] all_trainable = [ v for v in tf.trainable_variables() if 'beta' not in v.name and 'gamma' not in v.name ] fc_trainable = [v for v in all_trainable if 'fc' in v.name] conv_trainable = [v for v in all_trainable if 'fc' not in v.name] # 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]) raw_prediction100 = tf.reshape(raw_output100, [-1, args.num_classes]) raw_prediction075 = tf.reshape(raw_output075, [-1, args.num_classes]) raw_prediction05 = tf.reshape(raw_output05, [-1, args.num_classes]) #pdb.set_trace() 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] label_proc075 = prepare_label(label_batch, tf.stack(raw_output075.get_shape()[1:3]), num_classes=args.num_classes, one_hot=False) label_proc05 = prepare_label(label_batch, tf.stack(raw_output05.get_shape()[1:3]), num_classes=args.num_classes, one_hot=False) raw_gt = tf.reshape(label_proc, [ -1, ]) raw_gt075 = tf.reshape(label_proc075, [ -1, ]) raw_gt05 = tf.reshape(label_proc05, [ -1, ]) indices = tf.squeeze(tf.where(tf.less_equal(raw_gt, args.num_classes - 1)), 1) indices075 = tf.squeeze( tf.where(tf.less_equal(raw_gt075, args.num_classes - 1)), 1) indices05 = tf.squeeze( tf.where(tf.less_equal(raw_gt05, args.num_classes - 1)), 1) gt = tf.cast(tf.gather(raw_gt, indices), tf.int32) gt075 = tf.cast(tf.gather(raw_gt075, indices075), tf.int32) gt05 = tf.cast(tf.gather(raw_gt05, indices05), tf.int32) prediction = tf.gather(raw_prediction, indices) prediction100 = tf.gather(raw_prediction100, indices) prediction075 = tf.gather(raw_prediction075, indices075) prediction05 = tf.gather(raw_prediction05, indices05) # Pixel-wise softmax loss. loss = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=prediction, labels=gt) loss100 = tf.nn.sparse_softmax_cross_entropy_with_logits( logits=prediction100, labels=gt) loss075 = tf.nn.sparse_softmax_cross_entropy_with_logits( logits=prediction075, labels=gt075) loss05 = tf.nn.sparse_softmax_cross_entropy_with_logits( logits=prediction05, labels=gt05) 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.reduce_mean( loss100) + tf.reduce_mean(loss075) + tf.reduce_mean(loss05) + 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()) print('\n\n') print('-------------------------------------------------------') print('--------------- Trainable parameters ------------------') print('-------------------------------------------------------') total_params = 0 for v in tf.trainable_variables(): shape = v.get_shape() params = 1 for dim in shape: params *= dim.value print('{:<30s}: {:<20s}\t{:<10s}'.format(v.name, str(shape), str(params))) total_params += params print('total_pamars = {}'.format(total_params)) print('-------------------------------------------------------\n\n') # Define loss and optimisation parameters. 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)) 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) # Define a variable to accumulate gradients. accum_grads = [ tf.Variable(tf.zeros_like(v.initialized_value()), trainable=False) for v in conv_trainable + fc_w_trainable + fc_b_trainable ] # Define an operation to clear the accumulated gradients for next batch. zero_op = [v.assign(tf.zeros_like(v)) for v in accum_grads] # Compute gradients. grads = tf.gradients(reduced_loss, conv_trainable + fc_w_trainable + fc_b_trainable) # Accumulate and normalise the gradients. accum_grads_op = [ accum_grads[i].assign_add(grad / args.grad_update_every) for i, grad in enumerate(grads) ] grads_conv = accum_grads[:len(conv_trainable)] grads_fc_w = accum_grads[len(conv_trainable):(len(conv_trainable) + len(fc_w_trainable))] grads_fc_b = accum_grads[(len(conv_trainable) + len(fc_w_trainable)):] # Apply the gradients. 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 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) # Load variables if the checkpoint is provided. if args.restore_from is not None: loader = tf.train.Saver(var_list=restore_var) load(loader, sess, args.restore_from) # 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} loss_value = 0 # Clear the accumulated gradients. sess.run(zero_op, feed_dict=feed_dict) # Accumulate gradients. for i in range(args.grad_update_every): _, l_val = sess.run([accum_grads_op, reduced_loss], feed_dict=feed_dict) loss_value += l_val # Normalise the loss. loss_value /= args.grad_update_every # Apply gradients. if step % args.save_pred_every == 0: images, labels, summary, _ = sess.run( [image_batch, label_batch, total_summary, train_op], feed_dict=feed_dict) summary_writer.add_summary(summary, step) save(saver, sess, args.snapshot_dir, step) else: sess.run(train_op, feed_dict=feed_dict) duration = time.time() - start_time print('step {:d} \t loss = {:.3f}, ({:.3f} sec/step)'.format( step, loss_value, duration)) coord.request_stop() coord.join(threads)
def main(): """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) # Create queue coordinator. coord = tf.train.Coordinator() # Load reader. with tf.name_scope("create_inputs"): reader = ImageReader( args.data_dir, 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) # Create network. net = DeepLabResNetModel({'data': image_batch}, is_training=args.is_training, num_classes=args.num_classes) # For a small batch size, it is better to keep # the statistics of the BN layers (running means and variances) # frozen, and to not update the values provided by the pre-trained model. # If is_training=True, the statistics will be updated during the training. # Note that is_training=False still updates BN parameters gamma (scale) and beta (offset) # if they are presented in var_list of the optimiser definition. # Predictions. raw_output = net.layers['fc1_voc12'] # Which variables to load. Running means and variances are not trainable, # thus all_variables() should be restored. # Restore all variables, or all except the last ones. restore_var = [v for v in tf.global_variables() if 'fc' not in v.name or not args.not_restore_last] trainable = [v for v in tf.trainable_variables() if 'fc1_voc12' in v.name] # Fine-tune only the last layers. 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) gt = tf.reshape(label_proc, [-1, args.num_classes]) # Pixel-wise softmax loss. loss = tf.nn.softmax_cross_entropy_with_logits(logits=prediction, labels=gt) reduced_loss = tf.reduce_mean(loss) # Processed predictions. 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()) # Define loss and optimisation parameters. optimiser = tf.train.AdamOptimizer(learning_rate=args.learning_rate) optim = optimiser.minimize(reduced_loss, var_list=trainable) # Set up tf session and initialize variables. config = tf.ConfigProto() config.gpu_options.allow_growth = True 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=40) # Load variables if the checkpoint is provided. if args.restore_from is not None: loader = tf.train.Saver(var_list=restore_var) load(loader, sess, args.restore_from) # 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() if step % args.save_pred_every == 0: loss_value, images, labels, preds, summary, _ = sess.run([reduced_loss, image_batch, label_batch, pred, total_summary, optim]) summary_writer.add_summary(summary, step) save(saver, sess, args.snapshot_dir, step) else: loss_value, _ = sess.run([reduced_loss, optim]) duration = time.time() - start_time print('step {:d} \t loss = {:.3f}, ({:.3f} sec/step)'.format(step, loss_value, duration)) coord.request_stop() coord.join(threads)
def main(): """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) # Create queue coordinator. coord = tf.train.Coordinator() # Load reader. with tf.name_scope("create_inputs"): reader = ImageReader( args.data_dir, 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) # Create network. net = DeepLabResNetModelOri50gcnaspp({'data': image_batch}, is_training=args.is_training, num_classes=args.num_classes) # For a small batch size, it is better to keep # the statistics of the BN layers (running means and variances) # frozen, and to not update the values provided by the pre-trained model. # If is_training=True, the statistics will be updated during the training. # Note that is_training=False still updates BN parameters gamma (scale) and beta (offset) # if they are presented in var_list of the optimiser definition. # Predictions. raw_output = net.layers['fc1_voc12'] # Which variables to load. Running means and variances are not trainable, # thus all_variables() should be restored. restore_var = [v for v in tf.global_variables() if 'fc' not in v.name or not args.not_restore_last] all_trainable = [v for v in tf.trainable_variables() if 'beta' not in v.name and 'gamma' not in v.name] fc_trainable = [v for v in all_trainable if 'fc' in v.name] conv_trainable = [v for v in all_trainable if 'fc' not in v.name] # 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) 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) loss_summary = tf.summary.scalar('loss', reduced_loss) # 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. merged = tf.summary.merge_all() summary_writer = tf.summary.FileWriter(args.snapshot_dir, graph=tf.get_default_graph()) # Define loss and optimisation parameters. 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)) 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 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) # Load variables if the checkpoint is provided. if args.restore_from is not None: loader = tf.train.Saver(var_list=restore_var) load(loader, sess, args.restore_from) #ckpt = tf.train.get_checkpoint_state('./snapshots50/') #loader.restore(sess, ckpt.model_checkpoint_path) # 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, images, labels, preds, summary, _ = sess.run( [reduced_loss, image_batch, label_batch, pred, merged, train_op], feed_dict=feed_dict) summary_writer.add_summary(summary, step) save(saver, sess, args.snapshot_dir, step) else: loss_value, _ = sess.run([reduced_loss, train_op], feed_dict=feed_dict) duration = time.time() - start_time print('step {:d} \t loss = {:.3f}, ({:.3f} sec/step)'.format(step, loss_value, duration)) coord.request_stop() coord.join(threads)
def main(): with tf.device('/cpu:0'): args = get_arguments() with tf.name_scope('create_inputs'): image_batch = tf.placeholder(dtype=tf.float32, shape=[args.batch_size, 256, 256, 3]) label_batch = tf.placeholder(dtype=tf.uint8, shape=[args.batch_size, 256, 256, 1]) # 定义多进程读取数据 multi_process = Multiprocess_Data(data_path=args.data_list, batch_size=args.batch_size, capacity=100, num_threads=3) with tf.device('/gpu:0'): #create network net = DeepLabResNetModel({'data': image_batch}, is_training=True, num_classes=args.num_classes) raw_output = net.layers['fc1_voc12'] restore_var = tf.global_variables() trainable = tf.trainable_variables() with tf.device('/cpu:0'): for var in trainable: tf.summary.histogram(var.op.name, var) print('raw_outout:', raw_output) prediction = tf.reshape(raw_output, [-1, args.num_classes]) print('prediction:', prediction) with tf.device('/cpu:0'): label_proc = prepare_label(label_batch, tf.stack(raw_output.get_shape()[1:3]), num_classes=args.num_classes) print('label_proc:', label_batch) gt = tf.reshape(label_proc, [-1, args.num_classes]) print('gt:', gt) loss = tf.nn.softmax_cross_entropy_with_logits(logits=prediction, labels=gt) reduced_loss = tf.reduce_mean(loss) with tf.device('/cpu:0'): tf.summary.scalar('loss', reduced_loss) #Processed predictions raw_output_up = tf.image.resize_bilinear(raw_output, tf.shape(image_batch)[1:3, ]) with tf.device('/cpu:0'): input_batch = tf.squeeze(label_batch, squeeze_dims=[3]) label_one_hot = tf.one_hot(input_batch, depth=args.num_classes) with tf.device('/cpu:0'): iou_loss = tl.cost.iou_coe(raw_output_up, label_one_hot, axis=[1, 2, 3]) tf.summary.scalar('iou_loss', iou_loss) optimiser = tf.train.AdamOptimizer(learning_rate=args.learning_rate) optim = optimiser.minimize(reduced_loss, var_list=trainable) with tf.device('/cpu:0'): merged = tf.summary.merge_all() # Set up tf session and initialize variables. config = tf.ConfigProto() config.gpu_options.allow_growth = True sess = tf.Session(config=config) train_writer = tf.summary.FileWriter(tensorboard_path, sess.graph) init = tf.global_variables_initializer() sess.run(init) with tf.device('/cpu:0'): # Saver for storing checkpoints of the model. saver = tf.train.Saver(var_list=tf.global_variables(), max_to_keep=20) # Load variables if the checkpoint is provided. if args.restore_from is not None: loader = tf.train.Saver(var_list=restore_var) load(loader, sess, args.restore_from) start_epoch_time = 0 total_loss, total_iou, n_batch = 0, 0, 0 # Iterate over training steps. for step in range(args.num_steps): start_time = time.time() x_train, y_train = multi_process.shuffle_batch() feed_dict = {image_batch: x_train, label_batch: y_train} if (step + 1) % args.save_pred_every == 0: loss_value, IoU, _ = sess.run([reduced_loss, iou_loss, optim], feed_dict=feed_dict) total_loss += loss_value total_iou += IoU n_batch += 1 # save mode to args.snapshot_dir save(saver, sess, args.snapshot_dir, step) else: loss_value, IoU, _ = sess.run([reduced_loss, iou_loss, optim], feed_dict=feed_dict) total_loss += loss_value total_iou += IoU n_batch += 1 duration = time.time() - start_time print( 'step {:d} \t loss = {:.3f}, iou = {:.3f}, ({:.3f} sec/step)'. format(step, loss_value, IoU, duration)) if step % 100 == 0: summary_all = sess.run(merged, feed_dict=feed_dict) train_writer.add_summary(summary_all, step) print( " ** train_loss: %f iou: %f took %fs (2d with distortion)" % (total_loss / n_batch, total_iou / n_batch, time.time() - start_epoch_time)) total_loss, total_iou, n_batch = 0, 0, 0 start_epoch_time = time.time()