def main():
args = get_arguments()
cfg = Config(dataset="others",
is_training=False,
filter_scale=args.filter_scale)
cfg.param["num_classes"] = args.num_classes
net = ICNet_BN(cfg=cfg, mode='inference')
net.create_session()
net.restore(args.weights)
os.makedirs(args.output_dir, exist_ok=True)
input_files = os.listdir(args.input_dir)
input_files = [
f for f in input_files if ".png" in f or ".jpg" in f or ".jpeg" in f
]
for input_file in input_files:
input_path = os.path.join(args.input_dir, input_file)
img = cv2.imread(input_path)
# img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
if img.shape[0] == 400 and img.ndim == 3:
img = np.vstack(
[np.zeros(shape=(200, 800, 3), dtype=img.dtype), img])
if img.shape != cfg.INFER_SIZE:
img = cv2.resize(img, (cfg.INFER_SIZE[1], cfg.INFER_SIZE[0]))
prediction = net.predict(img)[0]
res = np.concatenate([img, prediction], axis=1)
output_path = os.path.join(args.output_dir, input_file)
cv2.imwrite(output_path, res)
print("Successfully ran inference on: " + input_file + " . Saved to " +
output_path)
def main():
"""Create the model and start the training."""
args = get_arguments()
"""
Get configurations here. We pass some arguments from command line to init configurations, for training hyperparameters,
you can set them in TrainConfig Class.
Note: we set filter scale to 1 for pruned model, 2 for non-pruned model. The filters numbers of non-pruned
model is two times larger than prunde model, e.g., [h, w, 64] <-> [h, w, 32].
"""
cfg = TrainConfig(dataset=args.dataset,
is_training=True,
random_scale=args.random_scale,
random_mirror=args.random_mirror,
filter_scale=args.filter_scale)
cfg.display()
# Setup training network and training samples
train_reader = ImageReader(cfg=cfg, mode='train')
train_net = ICNet_BN(image_reader=train_reader,
cfg=cfg, mode='train')
loss_sub4, loss_sub24, loss_sub124, reduced_loss = create_losses(train_net, train_net.labels, cfg)
# Setup validation network and validation samples
with tf.variable_scope('', reuse=True):
val_reader = ImageReader(cfg, mode='eval')
val_net = ICNet_BN(image_reader=val_reader,
cfg=cfg, mode='train')
val_loss_sub4, val_loss_sub24, val_loss_sub124, val_reduced_loss = create_losses(val_net, val_net.labels, cfg)
# Using Poly learning rate policy
base_lr = tf.constant(cfg.LEARNING_RATE)
step_ph = tf.placeholder(dtype=tf.float32, shape=())
learning_rate = tf.scalar_mul(base_lr, tf.pow((1 - step_ph / cfg.TRAINING_STEPS), cfg.POWER))
# Set restore variable
#restore_var = tf.global_variables()
restore_var = [v for v in tf.global_variables() ]#if 'conv6_cls' not in v.name]
all_trainable = [v for v in tf.trainable_variables() if ('beta' not in v.name and 'gamma' not in v.name) or args.train_beta_gamma]
# Gets moving_mean and moving_variance update operations from tf.GraphKeys.UPDATE_OPS
if args.update_mean_var == False:
update_ops = None
else:
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
#opt_conv = tf.train.MomentumOptimizer(learning_rate, cfg.MOMENTUM)
opt_conv = tf.train.AdamOptimizer(learning_rate=learning_rate)
grads = tf.gradients(reduced_loss, all_trainable)
train_op = opt_conv.apply_gradients(zip(grads, all_trainable))
# Create session & restore weights (Here we only need to use train_net to create session since we reuse it)
train_net.create_session()
saver = tf.train.Saver(var_list=restore_var, max_to_keep=10)
if cfg.model_weight is not None:
if cfg.model_weight.endswith(".npy"):
train_net.restore(cfg.model_weight, restore_var)
else:
saver.restore(train_net.sess, cfg.model_weight)
#train_net.restore(cfg.model_weight, all_trainable)
#saver = tf.train.Saver(var_list=tf.global_variables(), max_to_keep=5)
train_net.save(saver, cfg.SNAPSHOT_DIR, 0)
# Iterate over training steps.
for step in range(cfg.TRAINING_STEPS):
start_time = time.time()
feed_dict = {step_ph: step}
if step % cfg.SAVE_PRED_EVERY == 0:
loss_value, loss1, loss2, loss3, val_loss_value, _ = train_net.sess.run([reduced_loss, loss_sub4, loss_sub24, loss_sub124, val_reduced_loss, train_op], feed_dict=feed_dict)
train_net.save(saver, cfg.SNAPSHOT_DIR, step)
else:
loss_value, loss1, loss2, loss3, val_loss_value, _ = train_net.sess.run([reduced_loss, loss_sub4, loss_sub24, loss_sub124, val_reduced_loss, train_op], feed_dict=feed_dict)
duration = time.time() - start_time
print('step {:d} \t total loss = {:.3f}, sub4 = {:.3f}, sub24 = {:.3f}, sub124 = {:.3f}, val_loss: {:.3f} ({:.3f} sec/step)'. \
format(step, loss_value, loss1, loss2, loss3, val_loss_value, duration))
def main():
"""Create the model and start the training."""
args = get_arguments()
"""
Get configurations here. We pass some arguments from command line to init configurations, for training hyperparameters,
you can set them in TrainConfig Class.
Note: we set filter scale to 1 for pruned model, 2 for non-pruned model. The filters numbers of non-pruned
model is two times larger than prunde model, e.g., [h, w, 64] <-> [h, w, 32].
"""
cfg = TrainConfig(dataset=args.dataset,
is_training=True,
random_scale=args.random_scale,
random_mirror=args.random_mirror,
filter_scale=args.filter_scale)
if args.num_classes is not None:
cfg.param["num_classes"] = args.num_classes
if args.data_dir is not None:
cfg.param["data_dir"] = args.data_dir
if args.val_list is not None:
cfg.param["eval_list"] = args.val_list
if args.train_list is not None:
cfg.param["train_list"] = args.train_list
if args.ignore_label is not None:
cfg.param["ignore_label"] = args.ignore_label
if args.eval_size is not None:
cfg.param["eval_size"] = [
int(x.strip()) for x in args.eval_size.split("x")[::-1]
]
if args.training_size is not None:
cfg.TRAINING_SIZE = [
int(x.strip()) for x in args.training_size.split("x")[::-1]
]
if args.batch_size is not None:
cfg.BATCH_SIZE = args.batch_size
if args.learning_rate is not None:
cfg.LEARNING_RATE = args.learning_rate
if args.restore_from is not None:
cfg.model_weight = args.restore_from
if args.snapshot_dir is not None:
cfg.SNAPSHOT_DIR = args.snapshot_dir
if args.restore_from == "scratch":
from tqdm import tqdm
import cv2
import joblib as joblib
if not args.img_mean:
print(
"Calculating img mean for custom dataset. To prevent this, specify it with --img-mean next time"
)
image_files, annotation_files = read_labeled_image_list(
cfg.param["data_dir"], cfg.param["train_list"])
means = joblib.Parallel(n_jobs=6)(
joblib.delayed(calc_mean)(image_file, cv2)
for image_file in tqdm(image_files, desc="calc img mean"))
cfg.IMG_MEAN = np.mean(means, axis=0).tolist()
else:
cfg.IMG_MEAN = [float(x.strip()) for x in args.img_mean.split(",")]
cfg.display()
# Setup training network and training samples
train_reader = ImageReader(cfg=cfg, mode='train')
train_net = ICNet_BN(image_reader=train_reader, cfg=cfg, mode='train')
loss_sub4, loss_sub24, loss_sub124, reduced_loss = create_losses(
train_net, train_net.labels, cfg)
# Setup validation network and validation samples
with tf.variable_scope('', reuse=True):
val_reader = ImageReader(cfg, mode='eval')
val_net = ICNet_BN(image_reader=val_reader, cfg=cfg, mode='train')
val_loss_sub4, val_loss_sub24, val_loss_sub124, val_reduced_loss = create_losses(
val_net, val_net.labels, cfg)
# Using Poly learning rate policy
base_lr = tf.constant(cfg.LEARNING_RATE)
step_ph = tf.placeholder(dtype=tf.float32, shape=())
learning_rate = tf.scalar_mul(
base_lr, tf.pow((1 - step_ph / cfg.TRAINING_STEPS), cfg.POWER))
# Set restore variable
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)
or args.train_beta_gamma
]
# Gets moving_mean and moving_variance update operations from tf.GraphKeys.UPDATE_OPS
if args.update_mean_var == False:
update_ops = None
else:
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
opt_conv = tf.train.MomentumOptimizer(learning_rate, cfg.MOMENTUM)
grads = tf.gradients(reduced_loss, all_trainable)
train_op = opt_conv.apply_gradients(zip(grads, all_trainable))
# Create session & restore weights (Here we only need to use train_net to create session since we reuse it)
train_net.create_session()
if args.initializer:
train_net.set_initializer(initializer_algorithm=args.initializer)
train_net.initialize_variables()
if not args.restore_from or args.restore_from != "scratch":
train_net.restore(cfg.model_weight, restore_var)
saver = tf.train.Saver(var_list=tf.global_variables(), max_to_keep=20)
total_parameters = 0
for variable in tf.trainable_variables():
# shape is an array of tf.Dimension
shape = variable.get_shape()
variable_parameters = 1
for dim in shape:
variable_parameters *= dim.value
total_parameters += variable_parameters
print("Total trainable parameters: " + str(total_parameters))
# Iterate over training steps.
val_loss_value = 10.0
min_val_loss = float("inf")
stagnation = 0
max_non_decreasing_val_loss = int(
np.ceil(args.early_stopping_patience * len(train_reader.image_list) /
(cfg.BATCH_SIZE * cfg.EVAL_EVERY)))
print(
"Maximum times that val loss can stagnate before early stopping is applied: "
+ str(max_non_decreasing_val_loss))
for step in range(cfg.TRAINING_STEPS):
start_time = time.time()
feed_dict = {step_ph: step}
if step % cfg.EVAL_EVERY == 0:
loss_value, loss1, loss2, loss3, val_loss_value, _ = train_net.sess.run(
[
reduced_loss, loss_sub4, loss_sub24, loss_sub124,
val_reduced_loss, train_op
],
feed_dict=feed_dict)
if val_loss_value < min_val_loss:
print("New best val loss {:.3f}. Saving weights...".format(
val_loss_value))
train_net.save(
saver,
cfg.SNAPSHOT_DIR,
step,
model_name="val{:.3f}model.ckpt".format(val_loss_value))
min_val_loss = val_loss_value
stagnation = 0
else:
stagnation += 1
else:
loss_value, loss1, loss2, loss3, _ = train_net.sess.run(
[reduced_loss, loss_sub4, loss_sub24, loss_sub124, train_op],
feed_dict=feed_dict)
duration = time.time() - start_time
print('step {:d} \t total loss = {:.3f}, sub4 = {:.3f}, sub24 = {:.3f}, sub124 = {:.3f}, val_loss: {:.3f} ({:.3f} sec/step)'.\
format(step, loss_value, loss1, loss2, loss3, val_loss_value, duration))
if stagnation > max_non_decreasing_val_loss:
print("Early stopping")
break
def main():
"""Create the model and start the training."""
args = get_arguments()
"""
Get configurations here. We pass some arguments from command line to init configurations, for training hyperparameters,
you can set them in TrainConfig Class.
Note: we set filter scale to 1 for pruned model, 2 for non-pruned model. The filters numbers of non-pruned
model is two times larger than prunde model, e.g., [h, w, 64] <-> [h, w, 32].
"""
cfg = TrainConfig(dataset=args.dataset,
is_training=True,
random_scale=args.random_scale,
random_mirror=args.random_mirror,
filter_scale=args.filter_scale)
cfg.display()
# Setup training network and training samples
train_reader = ImageReader(cfg=cfg, mode='train')
train_net = ICNet_BN(image_reader=train_reader, cfg=cfg, mode='train')
loss_sub4, loss_sub24, loss_sub124, reduced_loss = create_losses(
train_net, train_net.labels, cfg)
# Setup validation network and validation samples
with tf.variable_scope('', reuse=True):
val_reader = ImageReader(cfg, mode='eval')
val_net = ICNet_BN(image_reader=val_reader, cfg=cfg, mode='train')
val_loss_sub4, val_loss_sub24, val_loss_sub124, val_reduced_loss = create_losses(
val_net, val_net.labels, cfg)
# Using Poly learning rate policy
base_lr = tf.constant(cfg.LEARNING_RATE)
step_ph = tf.placeholder(dtype=tf.float32, shape=())
learning_rate = tf.scalar_mul(
base_lr, tf.pow((1 - step_ph / cfg.TRAINING_STEPS), cfg.POWER))
# Set restore variable
# restore_var = tf.global_variables()
restore_var = [
v for v in tf.global_variables() if 'conv6_cls' not in v.name
]
all_trainable = [
v for v in tf.trainable_variables()
if ('beta' not in v.name and 'gamma' not in v.name)
or args.train_beta_gamma
]
# Gets moving_mean and moving_variance update operations from tf.GraphKeys.UPDATE_OPS
if args.update_mean_var == False:
update_ops = None
else:
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
opt_conv = tf.train.MomentumOptimizer(learning_rate, cfg.MOMENTUM)
grads = tf.gradients(reduced_loss, all_trainable)
train_op = opt_conv.apply_gradients(zip(grads, all_trainable))
# check_ckpt(cfg.model_weight)
# Create session & restore weights (Here we only need to use train_net to create session since we reuse it)
# tf.reset_default_graph()
train_net.create_session()
print("Restoring weights ...")
train_net.restore(cfg.model_weight, restore_var)
saver = tf.train.Saver(var_list=tf.global_variables(), max_to_keep=5)
# Losses for plotting
train_loss_list = []
val_loss_list = []
avg_train_loss_list = []
avg_val_loss_list = []
steps = []
fig, ax = plt.subplots(1, 1)
line1, = ax.plot(steps, train_loss_list, 'b--', linewidth=0.5)
line2, = ax.plot(steps, val_loss_list, 'g-', linewidth=0.5)
line1.set_label('Train Loss')
line2.set_label('Val Loss')
ax.legend()
plt.grid()
# Iterate over training steps.
for step in range(cfg.TRAINING_STEPS):
start_time = time.time()
feed_dict = {step_ph: step}
if step % cfg.SAVE_PRED_EVERY == 0:
loss_value, loss1, loss2, loss3, val_loss_value, _ = train_net.sess.run(
[
reduced_loss, loss_sub4, loss_sub24, loss_sub124,
val_reduced_loss, train_op
],
feed_dict=feed_dict)
train_net.save(saver, cfg.SNAPSHOT_DIR, step)
plt.savefig('snapshots/loss.png')
else:
loss_value, loss1, loss2, loss3, val_loss_value, _ = train_net.sess.run(
[
reduced_loss, loss_sub4, loss_sub24, loss_sub124,
val_reduced_loss, train_op
],
feed_dict=feed_dict)
duration = time.time() - start_time
print('step {:d} \t total loss = {:.3f}, sub4 = {:.3f}, sub24 = {:.3f}, sub124 = {:.3f}, val_loss: {:.3f} ({:.3f} sec/step)'.\
format(step, loss_value, loss1, loss2, loss3, val_loss_value, duration))
# Plot
train_loss_list.append(loss_value)
val_loss_list.append(val_loss_value)
if step % 50 == 0 and step != 0:
steps.append(step)
avg_train_loss_list.append(
sum(train_loss_list[-49:] + [loss_value]) / 50.0)
avg_val_loss_list.append(
sum(val_loss_list[-49:] + [val_loss_value]) / 50.0)
line1, = ax.plot(steps, avg_train_loss_list, 'b--', linewidth=0.5)
line2, = ax.plot(steps, avg_val_loss_list, 'g-', linewidth=0.5)
fig.canvas.draw()
plt.show(block=False)