def get_model(images, num_classes, for_eval=False):
net = PSPNet50(
{'data': images},
is_training=not for_eval,
num_classes=num_classes
)
return net
def main():
args = get_arguments()
if args.model == 'pspnet':
model = PSPNet50()
elif args.model == 'fcn':
model = FCN8s()
elif args.model == 'enet':
model = ENet()
elif args.model == 'icnet':
model = ICNet()
model.read_input(args.img_path)
# Init tf Session
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
init = tf.global_variables_initializer()
sess.run(init)
model.load(model_path[args.model], sess)
preds = model.forward(sess)
if not os.path.exists(args.save_dir):
os.makedirs(args.save_dir)
misc.imsave(args.save_dir + args.model + '_' + model.img_name, preds[0])
def main(args):
# device number
if args.gpu_num:
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
# set the necessary list
train_list = pd.read_csv(args.train_list, header=None)
val_list = pd.read_csv(args.val_list, header=None)
# set the necessary directories
trainimg_dir = args.trainimg_dir
trainmsk_dir = args.trainmsk_dir
valimg_dir = args.valimg_dir
valmsk_dir = args.valmsk_dir
# get old session old_session = KTF.get_session()
with tf.Graph().as_default():
session = tf.Session('')
KTF.set_session(session)
KTF.set_learning_phase(1)
# set callbacks
cp_cb = ModelCheckpoint(filepath=args.log_dir,
monitor='val_loss',
verbose=1,
save_best_only=True,
mode='auto',
period=2)
es_cb = EarlyStopping(monitor='val_loss',
patience=2,
verbose=1,
mode='auto')
tb_cb = TensorBoard(log_dir=args.log_dir, write_images=True)
# set generater
train_gen = data_gen_small(trainimg_dir, trainmsk_dir, train_list,
args.batch_size,
[args.input_shape[0], args.input_shape[1]],
args.n_labels)
val_gen = data_gen_small(valimg_dir, valmsk_dir, val_list,
args.batch_size,
[args.input_shape[0], args.input_shape[1]],
args.n_labels)
# set model
pspnet = PSPNet50(input_shape=args.input_shape,
n_labels=args.n_labels,
output_mode=args.output_mode,
upsample_type=args.upsample_type)
print(pspnet.summary())
# compile model
pspnet.compile(loss=args.loss,
optimizer=args.optimizer,
metrics=["accuracy"])
# fit with genarater
pspnet.fit_generator(generator=train_gen,
steps_per_epoch=args.epoch_steps,
epochs=args.n_epochs,
validation_data=val_gen,
validation_steps=args.val_steps,
callbacks=[cp_cb, es_cb, tb_cb])
def main():
"""
Usage: python test.py --model_name $Modelname$ --train_categories
gets val data and predicts with model $Modelname$
stores rgb sgmentations in logs/$Modelname$/pred
"""
parser = test_parser()
args = parser.parse_args()
if args.batch_size != 1:
raise ValueError("Evaluation is not defined for "
"batch size other than 1.")
# determine some variables from args
dims, input_shape, n_labels, model_dir, checkpoint_path = find_input(args)
weight_mask, weight_mask_iou = find_masks(args, n_labels)
# determine output directory
output_dir = args.output_dir + "/" + args.model_name
if args.clear_output_dir:
shutil.rmtree(output_dir, ignore_errors=True)
time.sleep(1)
os.makedirs(output_dir, exist_ok=True)
max_vis = 264
max_calc = VAL_STEPS
# load data gen
datagen = DataStorage(args.data_dir, args.dataset, dims, n_labels)
train_gen = datagen.generate(args.batch_size,
args.use_flow,
args.use_mb,
args.use_occ,
args.train_categories,
split="train",
shuffle=True,
normalize_images=not args.no_norm_images)
val_gen = datagen.generate(args.batch_size,
args.use_flow,
args.use_mb,
args.use_occ,
args.train_categories,
split="val",
shuffle=False,
random_crop=False,
write_folder=output_dir,
write_cropped=True,
write_uncropped=True,
write_max=max_vis,
normalize_images=not args.no_norm_images,
full_size=args.full_size)
# setup model
loss = weighted_categorical_crossentropy(weight_mask, args.weight_mult)
def mean_iou(y_true, y_pred):
return mean_iou_func(y_true,
y_pred,
args.batch_size,
dims,
n_labels,
weight_mask_iou,
iou_test=False)
# load and print model epoch
try:
model_epoch = int(open(model_dir + "/epochs.txt").read())
except OSError:
raise FileNotFoundError("could not find epochs.txt")
print("reloading model epoch", model_epoch)
if args.full_size:
# for full size models, load weights only instead of full model
# check weight age
weight_info = model_dir + "/epochs_weight.txt"
weight_path = model_dir + "/weights.h5"
if not os.path.exists(weight_info):
weight_epoch = -1
else:
weight_epoch = int(open(weight_info, "rt").read())
if weight_epoch != model_epoch:
print("weight epoch is", weight_epoch, "(-1 = no weights)")
print("load model and save weights...")
try:
pspnet = load_model(checkpoint_path,
custom_objects={
'mean_iou': mean_iou,
'loss': loss
})
except OSError:
print(checkpoint_path + " does not exist")
return
print("model loaded, saving weights...")
pspnet.save_weights(weight_path, overwrite=True)
open(weight_info, "wt").write(str(model_epoch))
# try to free up all that memory that the model consumed
del pspnet
KTF.clear_session()
# now weights should be correct. load them
channels = 3
if args.use_flow:
channels += 2
if args.use_mb:
channels += 1
if args.use_occ:
channels += 1
input_shape = (args.full_height, args.full_width, channels)
pspnet = PSPNet50(input_shape=input_shape,
n_labels=n_labels,
output_mode=args.output_mode,
upsample_type=args.upsample_type)
pspnet.compile(loss=loss,
optimizer=args.optimizer,
metrics=["accuracy", mean_iou])
pspnet.load_weights(weight_path)
print("loaded weights and created full size model")
else:
# regular size models
try:
pspnet = load_model(checkpoint_path,
custom_objects={
'mean_iou': mean_iou,
'loss': loss
})
except OSError:
print(checkpoint_path + " does not exist")
return
# check weights for nans
for i, l in enumerate(pspnet.layers):
ws = l.get_weights()
for j, w in enumerate(ws):
# print("weights shaped", w.shape)
nans = np.isnan(w)
nans_c = np.sum(nans)
if nans_c > 0:
print("layer", i)
print(l)
print("weights", j, "nans", nans_c)
os.makedirs(model_dir + '/pred', exist_ok=True)
color_mask = make_color_mask(n_labels, args.train_categories)
print("saving first", max_vis, "inputs to",
output_dir + "/" + args.model_name)
print("categories?", args.train_categories)
evalus = []
verbose = 1
preds, targets = [], []
for i in range(VAL_STEPS):
# print progress
if i % 10 == 0:
print("{}/{} ({:.1f}%)".format(i, VAL_STEPS, i / VAL_STEPS * 100))
# get data
x, target = next(val_gen)
print("x", x.shape, "y", target.shape)
continue
# predict and save evaluation (loss etc.)
prediction = pspnet.predict(x, batch_size=1, verbose=verbose)
evalu = pspnet.evaluate(x, target, batch_size=args.batch_size)
evalus.append(evalu)
# only calculate on the first Y images then break
if i >= max_calc:
break
# only visualize the first X images
if i >= max_vis:
verbose = 0
continue
print("PRED", prediction)
# reshape prediction and target
if args.full_size:
h = args.full_height
w = args.full_width
else:
h = args.img_height
w = args.img_width
prediction = np.reshape(prediction, (args.batch_size, h, w, n_labels))
target = np.reshape(target, (args.batch_size, h, w, n_labels))
# color prediction and target
img = decode_onehot(prediction)
mask = decode_onehot(target)
preds.append(img)
targets.append(mask)
img_new = apply_color_mask(img, color_mask)
mask_new = apply_color_mask(mask, color_mask)
# save prediction and target
full_str = ""
if args.full_size:
full_str = "full"
img_file = model_dir + '/pred/pred' + full_str + str(i) + '.png'
mask_file = model_dir + '/pred/pred' + full_str + str(i) + 'true.png'
write_image(img_file, img_new)
write_image(mask_file, mask_new)
print("saved prediction", i, "to", img_file)
# evaluate metrics
names = ["val loss", "val acc", "val mean iou"]
evalus = np.array(evalus)
evalus_mean = np.mean(evalus, axis=1)
print()
print("averages over each individual prediction (wrong values)")
for i, n in enumerate(names):
print("{:20} {:.4f}".format(n, evalus_mean[i]))
# manually compute iou
targets = np.array(targets)
preds = np.array(preds)
print()
print("mean numpy iou over everything with weights")
iou, seen_c = compute_mean_iou(targets,
preds,
n_labels,
weights=weight_mask_iou,
return_seen=True)
print("{:20} {:.4f} seen {:2d} classes".format("numpy val mean iou", iou,
seen_c))
print()
print("class predictions numpy iou")
ious = []
for i in range(n_labels):
test_weight_mask = np.zeros(n_labels)
test_weight_mask[i] = 1
iou, seen_c = compute_mean_iou(targets,
preds,
n_labels,
weights=test_weight_mask,
return_seen=True,
divide_over_num_weights=True)
print("iou class {:3d}: {:.4f} seen {:2.0f} classes".format(
i, iou, seen_c))
ious.append(iou)
print("final iou {:4f}".format(np.mean(ious)))
print()
csv_file = model_dir + "/training.log"
print("last entry from training log")
fh = open(csv_file, "rt")
csv_reader = csv.reader(fh)
title_row = next(csv_reader)
row = None
for row in csv_reader:
pass
for title, val in zip(title_row, row):
print("{:20} {:.4f}".format(title, float(val)))
print()
print("test it with evaluate_generator")
evalus2 = pspnet.evaluate_generator(val_gen, VAL_STEPS, max_queue_size=1)
for i, n in enumerate(names):
print("{:20} {:.4f}".format(n, evalus2[i]))
print("done testing.")
def main():
parser = train_parser()
args = parser.parse_args()
if args.use_cpu:
# disable gpu completely
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
else:
# set device number
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
if args.seed != -1:
# set random seed
np.random.seed(args.seed)
random.seed(args.seed)
tf.set_random_seed(args.seed)
# get old session old_session = KTF.get_session()
with tf.Graph().as_default():
session = tf.Session('')
KTF.set_session(session)
KTF.set_learning_phase(1)
# determine some variables from args
dims, input_shape, n_labels, model_dir, checkpoint_path = find_input(
args)
os.makedirs(model_dir, exist_ok=True)
# set callbacks
# for multiple checkpoints set filepath to "...point{epoch:d}.hdf5"
cp_cb = MyCheckpoint(filepath=model_dir + "/checkpoint.hdf5",
monitor='val_loss',
verbose=1,
save_best_only=False,
mode='auto',
period=1)
tb_cb = TensorBoard(log_dir=model_dir, write_images=True)
log = CSVLogger(model_dir + "/training.log", append=True)
# set generator
if args.use_random:
# create random data
train_gen = data_gen_random(dims, n_labels, args.batch_size,
args.use_flow)
val_gen = data_gen_random(dims, n_labels, args.batch_size,
args.use_flow)
elif args.use_numpy:
# load small numpy dataset
n_labels = args.n_classes
datagen = DataStorageNpy(args.train_imgs_npy, args.train_mask_npy,
args.train_flow_npy, args.val_percent,
dims, n_labels, args.use_flow,
args.batch_size)
train_gen = datagen.generate(training=True)
val_gen = datagen.generate(training=False)
else:
# load dataset from folders
datagen = DataStorage(args.data_dir, args.dataset, dims, n_labels)
train_gen = datagen.generate(
args.batch_size,
args.use_flow,
args.use_mb,
args.use_occ,
args.train_categories,
split="train",
shuffle=True,
normalize_images=not args.no_norm_images)
val_gen = datagen.generate(
args.batch_size,
args.use_flow,
args.use_mb,
args.use_occ,
args.train_categories,
split="val",
shuffle=False,
random_crop=False,
normalize_images=not args.no_norm_images)
if args.test_generator:
imgs, labels = next(train_gen)
print("datagen yielded", imgs.shape, labels.shape)
return
# # test has flow and images, but masks are zero
# # so it cannot be used to compute accuracy
# test_gen = datagen.generate(args.batch_size, args.use_flow,
# args.use_mb, args.use_occ,
# split="test")
# ----- determine class weights
weight_mask, weight_mask_iou = find_masks(args, n_labels)
print("using weight mask:", weight_mask)
print("using weight mask for iou:", weight_mask_iou)
# ----- define loss
old_loss = args.loss
print("old loss", old_loss)
loss = weighted_categorical_crossentropy(weight_mask, args.weight_mult)
# ----- define iou metric
def mean_iou(y_true, y_pred):
return mean_iou_func(y_true,
y_pred,
args.batch_size,
dims,
n_labels,
weight_mask_iou,
iou_test=False)
# ----- restart or load model
restart = args.restart
if not args.restart:
# check if model is available
if not os.path.exists(checkpoint_path):
print("no previous model available, restarting from epoch 0")
restart = True
print("model input shape", input_shape)
if restart:
# set model
pspnet = PSPNet50(input_shape=input_shape,
n_labels=n_labels,
output_mode=args.output_mode,
upsample_type=args.upsample_type)
# compile model
pspnet.compile(loss=loss,
optimizer=args.optimizer,
metrics=["accuracy", mean_iou])
# metrics=["acc_iou", mean_iou])
starting_epoch = 0
else:
# load model from dir
try:
pspnet = load_model(checkpoint_path,
custom_objects={
'mean_iou': mean_iou,
'loss': loss
})
except OSError:
raise OSError("failed loading checkpoint", checkpoint_path)
except ValueError:
raise ValueError(
"possible the checkpoint file is corrupt because the "
"script died while saving. use the backup old_checkpoint")
try:
with open(model_dir + "/epochs.txt", "r") as fh:
starting_epoch = int(fh.read())
except OSError:
raise FileNotFoundError("could not find epochs.txt")
print("reloading model epoch", starting_epoch)
# print model summary (weights, layers etc)
if args.summary:
print(pspnet.summary())
# ----- fit with genarater
pspnet.fit_generator(generator=train_gen,
steps_per_epoch=args.epoch_steps,
epochs=args.n_epochs,
validation_data=val_gen,
validation_steps=args.val_steps,
callbacks=[cp_cb, tb_cb, log],
initial_epoch=starting_epoch)
def main():
"""Predict with ADE20K PSPNet50."""
parser = argparse.ArgumentParser()
parser.add_argument("weights_path", help="Path to ade20k h5 weights")
parser.add_argument("input_image",
help="Path to image we should predict on")
args = parser.parse_args()
# We're not training so provided learning_rate doesn't matter
net = PSPNet50(SEMSEG_SHAPE,
"pspnet50_ade20k",
SEMSEG_NUM_CLASSES,
learning_rate=1e-4)
with net.graph.as_default():
output_dir = os.path.dirname(args.input_image)
input_fname = os.path.splitext(os.path.basename(args.input_image))[0]
# Read the input image and perform ADE20K preprcoessing
img = misc.imread(args.input_image, mode='RGB')
h_ori, w_ori = img.shape[:2]
img = misc.imresize(img, SEMSEG_SHAPE)
img = img - DATA_MEAN
img = img[:, :, ::-1] # RGB => BGR
img = img.astype('float32')
# Load the weights from h5 file and assign them to corresponding tf tensors
weights = h5py.File(args.weights_path, 'r')
for i in tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES):
layer_name, tensor_name = i.name.split('/')[-2], i.name.split(
'/')[-1]
t = net.graph.get_tensor_by_name(i.name)
load_weights = weights[layer_name][layer_name][tensor_name].value
if tensor_name in [
"gamma:0", "beta:0", "moving_mean:0", "moving_variance:0"
]:
load_weights = load_weights.reshape(-1)
net.sess.run(tf.assign(t, load_weights))
# Predict on the image
feed_dict = {
"input_image:0": np.expand_dims(img, axis=0),
"is_training:0": False,
"keep_probability:0": 1.0
}
pred = net.sess.run(net.logits, feed_dict=feed_dict)[0]
# If predict shape isn't the same as original shape, make predict shape the same
if img.shape[0:1] != SEMSEG_SHAPE[0]:
h, w = pred.shape[:2]
pred = ndimage.zoom(pred, (1. * h_ori / h, 1. * w_ori / w, 1.),
order=1,
prefilter=False)
cm = np.argmax(pred, axis=2)
# Create human digestable prediction images
color_cm = add_color(cm)
alpha_blended = 0.5 * color_cm * 255 + 0.5 * misc.imread(
args.input_image, mode='RGB')
color_path = os.path.join(output_dir,
"{}_colored_pred.jpg".format(input_fname))
blend_path = os.path.join(output_dir,
"{}_blended_pred.jpg".format(input_fname))
misc.imsave(color_path, color_cm)
misc.imsave(blend_path, alpha_blended)
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)
coord = tf.train.Coordinator()
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)
net = PSPNet50({'data': image_batch},
is_training=True,
num_classes=args.num_classes)
raw_output = net.layers['conv6']
# According from the prototxt in Caffe implement, learning rate must multiply by 10.0 in pyramid module
fc_list = [
'conv5_3_pool1_conv', 'conv5_3_pool2_conv', 'conv5_3_pool3_conv',
'conv5_3_pool6_conv', 'conv6', 'conv5_4'
]
restore_var = [v for v in tf.global_variables()]
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
]
fc_trainable = [
v for v in all_trainable if v.name.split('/')[0] in fc_list
]
conv_trainable = [
v for v in all_trainable if v.name.split('/')[0] not in fc_list
] # lr * 1.0
fc_w_trainable = [v for v in fc_trainable
if 'weights' in v.name] # lr * 10.0
fc_b_trainable = [v for v in fc_trainable
if 'biases' in v.name] # lr * 20.0
assert (len(all_trainable) == len(fc_trainable) + len(conv_trainable))
assert (len(fc_trainable) == len(fc_w_trainable) + len(fc_b_trainable))
# Predictions: ignoring all predictions with labels greater or equal than n_classes
raw_prediction = tf.reshape(raw_output, [-1, args.num_classes])
label_proc = prepare_label(label_batch,
tf.stack(raw_output.get_shape()[1:3]),
num_classes=args.num_classes,
one_hot=False) # [batch_size, h, w]
raw_gt = tf.reshape(label_proc, [
-1,
])
indices = tf.squeeze(tf.where(tf.less_equal(raw_gt, args.num_classes - 1)),
1)
gt = tf.cast(tf.gather(raw_gt, indices), tf.int32)
prediction = tf.gather(raw_prediction, indices)
# Pixel-wise softmax loss.
loss = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=prediction,
labels=gt)
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)
# Using Poly learning rate policy
base_lr = tf.constant(args.learning_rate)
step_ph = tf.placeholder(dtype=tf.float32, shape=())
learning_rate = tf.scalar_mul(
base_lr, tf.pow((1 - step_ph / args.num_steps), args.power))
# Gets moving_mean and moving_variance update operations from tf.GraphKeys.UPDATE_OPS
if args.update_mean_var == False:
update_ops = None
else:
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
opt_conv = tf.train.MomentumOptimizer(learning_rate, args.momentum)
opt_fc_w = tf.train.MomentumOptimizer(learning_rate * 10.0,
args.momentum)
opt_fc_b = tf.train.MomentumOptimizer(learning_rate * 20.0,
args.momentum)
grads = tf.gradients(reduced_loss,
conv_trainable + fc_w_trainable + fc_b_trainable)
grads_conv = grads[:len(conv_trainable)]
grads_fc_w = grads[len(conv_trainable):(len(conv_trainable) +
len(fc_w_trainable))]
grads_fc_b = grads[(len(conv_trainable) + len(fc_w_trainable)):]
train_op_conv = opt_conv.apply_gradients(
zip(grads_conv, conv_trainable))
train_op_fc_w = opt_fc_w.apply_gradients(
zip(grads_fc_w, fc_w_trainable))
train_op_fc_b = opt_fc_b.apply_gradients(
zip(grads_fc_b, fc_b_trainable))
train_op = tf.group(train_op_conv, train_op_fc_w, train_op_fc_b)
# Set up tf session and initialize variables.
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
init = tf.global_variables_initializer()
sess.run(init)
# Saver for storing checkpoints of the model.
saver = tf.train.Saver(var_list=tf.global_variables(), max_to_keep=10)
ckpt = tf.train.get_checkpoint_state(SNAPSHOT_DIR)
#if ckpt and ckpt.model_checkpoint_path:
loader = tf.train.Saver(var_list=restore_var)
#load_step = int(os.path.basename(ckpt.model_checkpoint_path).split('-')[1])
load(loader, sess, RESTORE_FROM) #ckpt.model_checkpoint_path)
# else:
# print('No checkpoint file found.')
# load_step = 0
# Start queue threads.
threads = tf.train.start_queue_runners(coord=coord, sess=sess)
# Iterate over training steps.
for step in range(args.num_steps):
start_time = time.time()
feed_dict = {step_ph: step}
if step % args.save_pred_every == 0:
loss_value, _ = sess.run([reduced_loss, train_op],
feed_dict=feed_dict)
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)
coord = tf.train.Coordinator()
with tf.name_scope("create_inputs"):
reader = ImageReader(args.data_list, input_size, args.random_scale,
args.random_mirror, args.ignore_label, IMG_MEAN,
coord)
image_batch, label_batch = reader.dequeue(args.batch_size)
net = PSPNet50({'data': image_batch},
is_training=True,
num_classes=args.num_classes)
raw_output = net.layers['conv6']
# According from the prototxt in Caffe implement, learning rate must multiply by 10.0 in pyramid module
fc_list = [
'conv5_3_pool1_conv', 'conv5_3_pool2_conv', 'conv5_3_pool3_conv',
'conv5_3_pool6_conv', 'conv6', 'conv5_4'
]
restore_var = [
v for v in tf.global_variables()
if not (len([f for f in fc_list
if f in v.name])) or not args.not_restore_last
]
all_trainable = [
v for v in tf.trainable_variables()
if 'gamma' not in v.name and 'beta' not in v.name
]
fc_trainable = [
v for v in all_trainable if v.name.split('/')[0] in fc_list
]
conv_trainable = [
v for v in all_trainable if v.name.split('/')[0] not in fc_list
] # lr * 1.0
fc_w_trainable = [v for v in fc_trainable
if 'weights' in v.name] # lr * 10.0
fc_b_trainable = [v for v in fc_trainable
if 'biases' in v.name] # lr * 20.0
assert (len(all_trainable) == len(fc_trainable) + len(conv_trainable))
assert (len(fc_trainable) == len(fc_w_trainable) + len(fc_b_trainable))
# Predictions: ignoring all predictions with labels greater or equal than n_classes
raw_prediction = tf.reshape(raw_output, [-1, args.num_classes])
label_proc = prepare_label(label_batch,
tf.stack(raw_output.get_shape()[1:3]),
num_classes=args.num_classes,
one_hot=False) # [batch_size, h, w]
raw_gt = tf.reshape(label_proc, [
-1,
])
indices = tf.squeeze(tf.where(tf.less_equal(raw_gt, args.num_classes - 1)),
1)
gt = tf.cast(tf.gather(raw_gt, indices), tf.int32)
prediction = tf.gather(raw_prediction, indices)
# Pixel-wise softmax loss.
loss = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=prediction,
labels=gt)
# Make mistakes for class N more important for network
if USE_CLASS_WEIGHTS:
if len(CLASS_WEIGHTS) != NUM_CLASSES:
print('Incorrect class weights, it will be not used')
else:
mask = tf.zeros_like(loss)
for i, w in enumerate(CLASS_WEIGHTS):
mask = mask + tf.cast(tf.equal(gt, i),
tf.float32) * tf.constant(w)
loss = loss * mask
l2_losses = [
args.weight_decay * tf.nn.l2_loss(v) for v in tf.trainable_variables()
if 'weights' in v.name
]
reduced_loss = tf.reduce_mean(loss) + tf.add_n(l2_losses)
# Processed predictions: for visualisation.
raw_output_up = tf.image.resize_bilinear(raw_output,
tf.shape(image_batch)[1:3, ])
raw_output_up = tf.argmax(raw_output_up, dimension=3)
pred = tf.expand_dims(raw_output_up, dim=3)
# Image summary.
images_summary = tf.py_func(inv_preprocess,
[image_batch, args.save_num_images, IMG_MEAN],
tf.uint8)
labels_summary = tf.py_func(
decode_labels, [label_batch, args.save_num_images, args.num_classes],
tf.uint8)
preds_summary = tf.py_func(decode_labels,
[pred, args.save_num_images, args.num_classes],
tf.uint8)
total_summary = tf.summary.image(
'images',
tf.concat(axis=2,
values=[images_summary, labels_summary, preds_summary]),
max_outputs=args.save_num_images) # Concatenate row-wise.
summary_writer = tf.summary.FileWriter(args.snapshot_dir,
graph=tf.get_default_graph())
# Using Poly learning rate policy
base_lr = tf.constant(args.learning_rate)
step_ph = tf.placeholder(dtype=tf.float32, shape=())
learning_rate = tf.scalar_mul(
base_lr, tf.pow((1 - step_ph / args.num_steps), args.power))
# Gets moving_mean and moving_variance update operations from tf.GraphKeys.UPDATE_OPS
if args.update_mean_var == False:
update_ops = None
else:
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
opt_conv = tf.train.MomentumOptimizer(learning_rate, args.momentum)
opt_fc_w = tf.train.MomentumOptimizer(learning_rate * 10.0,
args.momentum)
opt_fc_b = tf.train.MomentumOptimizer(learning_rate * 20.0,
args.momentum)
grads = tf.gradients(reduced_loss,
conv_trainable + fc_w_trainable + fc_b_trainable)
grads_conv = grads[:len(conv_trainable)]
grads_fc_w = grads[len(conv_trainable):(len(conv_trainable) +
len(fc_w_trainable))]
grads_fc_b = grads[(len(conv_trainable) + len(fc_w_trainable)):]
train_op_conv = opt_conv.apply_gradients(
zip(grads_conv, conv_trainable))
train_op_fc_w = opt_fc_w.apply_gradients(
zip(grads_fc_w, fc_w_trainable))
train_op_fc_b = opt_fc_b.apply_gradients(
zip(grads_fc_b, fc_b_trainable))
train_op = tf.group(train_op_conv, train_op_fc_w, train_op_fc_b)
# Set up tf session and initialize variables.
config = tf.ConfigProto()
# config.gpu_options.allow_growth = True
# config.allow_soft_placement = True
# config.intra_op_parallelism_threads = 1
sess = tf.Session(config=config)
init = tf.global_variables_initializer()
sess.run(init)
# Saver for storing checkpoints of the model.
saver = tf.train.Saver(var_list=tf.global_variables(), max_to_keep=10)
ckpt = tf.train.get_checkpoint_state(SNAPSHOT_DIR)
if ckpt and ckpt.model_checkpoint_path:
loader = tf.train.Saver(var_list=restore_var)
load_step = int(
os.path.basename(ckpt.model_checkpoint_path).split('-')[1])
load(loader, sess, ckpt.model_checkpoint_path)
else:
print('No checkpoint file found.')
load_step = 0
# Start queue threads.
threads = tf.train.start_queue_runners(coord=coord, sess=sess)
# Iterate over training steps.
for step in range(args.num_steps):
start_time = time.time()
feed_dict = {step_ph: step}
if step % args.save_pred_every == 0:
loss_value, _, summary = sess.run(
[reduced_loss, train_op, total_summary], feed_dict=feed_dict)
summary_writer.add_summary(summary, step)
save(saver, sess, args.snapshot_dir, step)
else:
z, t, o, p, loss_value, _ = sess.run(
[raw_gt, raw_output, gt, prediction, reduced_loss, train_op],
feed_dict=feed_dict)
print(z.shape, t.shape, o.shape, p.shape)
duration = time.time() - start_time
print('step {:d} \t loss = {:.3f}, ({:.3f} sec/step)'.format(
step, loss_value, duration))
coord.request_stop()
coord.join(threads)
KTF.set_learning_phase(1)
# ### end
#data and labels
training_data_list = os.listdir(training_dataset_dir)
train_list = []
for data in training_data_list :
train_list.append(data.split(".")[0])
#data and labels
validation_data_list = os.listdir(validation_dataset_dir)
validation_list = []
for data in validation_data_list :
validation_list.append(data.split(".")[0])
model = PSPNet50(input_shape=input_shape, NUM_OF_CLASSES=NUM_OF_CLASSES)
model.compile(loss=crossentropy, optimizer='adagrad', metrics=["accuracy"])
# save model's architecture
model_json_str = model.to_json()
### add for TensorBoard
tb_cb = keras.callbacks.TensorBoard(log_dir=log_dir, histogram_freq=0, write_graph=True, write_images=True)
### end
#figure model's summary
#model.summary()