def get_init_fn(self, checkpoint_path):
"""Returns a function run by the chief worker to warm-start the training."""
checkpoint_exclude_scopes=["InceptionV4/Logits", "InceptionV4/AuxLogits"]
exclusions = [scope.strip() for scope in checkpoint_exclude_scopes]
variables_to_restore = []
for var in slim.get_model_variables():
excluded = False
for exclusion in exclusions:
if var.op.name.startswith(exclusion):
excluded = True
break
if not excluded:
variables_to_restore.append(var)
return slim.assign_from_checkpoint_fn(
checkpoint_path,
variables_to_restore)
python类assign_from_checkpoint_fn()的实例源码
def train(self):
s = tf.Session()
init_fn = slim.assign_from_checkpoint_fn("./vgg_19.ckpt", slim.get_variables_to_restore(exclude = ['generate_image']))
#optimizer = tf.train.AdamOptimizer(learning_rate = 1e-1, beta1 = 0.5, beta2 = 0.5).minimize(self.loss, var_list = [self.target])
optimizer = tf.contrib.opt.ScipyOptimizerInterface(self.loss, options={'maxiter': 1000}, var_list = [self.target])
s.run(tf.global_variables_initializer())
init_fn(s)
#for i in range(10000):
# _, loss_out = s.run([optimizer, self.loss])
# print("Current loss is: %.3f" %loss_out, end="\r")
#print("")
optimizer.minimize(s)
loss_out = s.run(self.loss)
print("Final loss: %.3f" %loss_out)
plt.imshow(np.clip(s.run(self.target), 0, 255).astype(np.uint8))
plt.show()
def create_init_fn_to_restore(master_checkpoint, train_dir):
"""Creates an init operations to restore weights from various checkpoints.
master_checkpoint is path to a checkpoint which contains all weights for
the whole model.
"""
if master_checkpoint is None:
return None
# Warn the user if a checkpoint exists in the train_dir. Then we'll be
# ignoring the checkpoint path anyway.
if tf.train.latest_checkpoint(train_dir):
tf.logging.info(
'Ignoring --checkpoint_path because a checkpoint already exists in %s'
% train_dir)
return None
if tf.gfile.IsDirectory(master_checkpoint):
checkpoint_path = tf.train.latest_checkpoint(master_checkpoint)
else:
checkpoint_path = master_checkpoint
tf.logging.info('Fine-tuning from %s' % checkpoint_path)
return slim.assign_from_checkpoint_fn(checkpoint_path, slim.get_model_variables())
def load_model_from_checkpoint_fn(self, model_fn):
"""Load weights from file and keep in memory.
Args:
model_fn: saved model file.
"""
# self.dm_model.use_graph()
print "start loading from checkpoint file..."
if self.vars_to_restore is None:
self.vars_to_restore = slim.get_variables()
restore_fn = slim.assign_from_checkpoint_fn(model_fn, self.vars_to_restore)
print "restoring model from {}".format(model_fn)
restore_fn(self.sess)
print "model restored."
def load_model(ckpt_dir, variables_to_restore=None):
"""Load model weights.
Assuming the model graph has been built.
Args:
ckpt_dir: checkpoint directory.
variables_to_restore: which variables to load from checkpoint.
"""
if not os.path.exists(ckpt_dir):
print "checkpoint dir {} not exist.".format(ckpt_dir)
return
ckpts = glob.glob(os.path.join(ckpt_dir, "*.ckpt*"))
ckpt = ckpts[0]
if variables_to_restore is None:
saver = tf.train.Saver()
else:
saver = tf.train.Saver(variables_to_restore)
with tf.Session() as sess:
# ckpt = tf.train.latest_checkpoint(ckpt_dir)
if ckpt:
saver.restore(sess, ckpt)
print "model loaded from {}".format(ckpt)
else:
print "unable to load model from {}".format(ckpt)
# another way.
# slim.assign_from_checkpoint_fn(
# ckpt,
# variables_to_restore,
# ignore_missing_vars=True)
def use_inceptionv4(self):
image_size = inception.inception_v4.default_image_size
img_path = "../../data/misec_images/EnglishCockerSpaniel_simon.jpg"
checkpoint_path = "../../data/trained_models/inception_v4/inception_v4.ckpt"
with tf.Graph().as_default():
image_string = tf.read_file(img_path)
image = tf.image.decode_jpeg(image_string, channels=3)
processed_image = inception_preprocessing.preprocess_image(image, image_size, image_size, is_training=False)
processed_images = tf.expand_dims(processed_image, 0)
# Create the model, use the default arg scope to configure the batch norm parameters.
with slim.arg_scope(inception.inception_v4_arg_scope()):
logits, _ = inception.inception_v4(processed_images, num_classes=1001, is_training=False)
probabilities = tf.nn.softmax(logits)
init_fn = slim.assign_from_checkpoint_fn(
checkpoint_path,
slim.get_model_variables('InceptionV4'))
with tf.Session() as sess:
init_fn(sess)
np_image, probabilities = sess.run([image, probabilities])
probabilities = probabilities[0, 0:]
sorted_inds = [i[0] for i in sorted(enumerate(-probabilities), key=lambda x:x[1])]
self.disp_names(sorted_inds,probabilities)
plt.figure()
plt.imshow(np_image.astype(np.uint8))
plt.axis('off')
plt.title(img_path)
plt.show()
return
def use_vgg16(self):
with tf.Graph().as_default():
image_size = vgg.vgg_16.default_image_size
img_path = "../../data/misec_images/First_Student_IC_school_bus_202076.jpg"
checkpoint_path = "../../data/trained_models/vgg16/vgg_16.ckpt"
image_string = tf.read_file(img_path)
image = tf.image.decode_jpeg(image_string, channels=3)
processed_image = vgg_preprocessing.preprocess_image(image, image_size, image_size, is_training=False)
processed_images = tf.expand_dims(processed_image, 0)
# Create the model, use the default arg scope to configure the batch norm parameters.
with slim.arg_scope(vgg.vgg_arg_scope()):
# 1000 classes instead of 1001.
logits, _ = vgg.vgg_16(processed_images, num_classes=1000, is_training=False)
probabilities = tf.nn.softmax(logits)
init_fn = slim.assign_from_checkpoint_fn(
checkpoint_path,
slim.get_model_variables('vgg_16'))
with tf.Session() as sess:
init_fn(sess)
np_image, probabilities = sess.run([image, probabilities])
probabilities = probabilities[0, 0:]
sorted_inds = [i[0] for i in sorted(enumerate(-probabilities), key=lambda x:x[1])]
self.disp_names(sorted_inds,probabilities,include_background=False)
plt.figure()
plt.imshow(np_image.astype(np.uint8))
plt.axis('off')
plt.title(img_path)
plt.show()
return
def create_resnet_model(img_dim):
pre_image = tf.placeholder(tf.float32, [None, None, 3])
processed_image = cnn_preprocessing.preprocess_for_eval(pre_image/255.0, img_dim, img_dim)
images = tf.placeholder(tf.float32, [None, img_dim, img_dim, 3])
# mean = tf.constant([123.68, 116.779, 103.939], dtype=tf.float32, shape=[1, 1, 1, 3], name='img_mean')
# processed_images = images - mean
with slim.arg_scope(resnet_utils.resnet_arg_scope()):
probs, endpoints = resnet_v2.resnet_v2_152(images, num_classes=1001, is_training = False)
print endpoints['resnet_v2_152/block4']
init_fn = slim.assign_from_checkpoint_fn(
'Data/CNNModels/resnet_v2_152.ckpt',
slim.get_model_variables('resnet_v2_152'))
sess = tf.Session()
init_fn(sess)
return {
'images_placeholder' : images,
'block4' : endpoints['resnet_v2_152/block4'],
'session' : sess,
'processed_image' : processed_image,
'pre_image' : pre_image,
'probs' : probs
}
def main():
model = config.get('config', 'model')
yolo = importlib.import_module('model.' + model)
width = config.getint(model, 'width')
height = config.getint(model, 'height')
with tf.Session() as sess:
image = tf.placeholder(tf.float32, [1, height, width, 3], name='image')
builder = yolo.Builder(args, config)
builder(image)
global_step = tf.contrib.framework.get_or_create_global_step()
model_path = tf.train.latest_checkpoint(utils.get_logdir(config))
tf.logging.info('load ' + model_path)
slim.assign_from_checkpoint_fn(model_path, tf.global_variables())(sess)
tf.logging.info('global_step=%d' % sess.run(global_step))
path = os.path.expanduser(os.path.expandvars(args.path))
if os.path.isfile(path):
detect(sess, builder.model, builder.names, image, path)
plt.show()
else:
for dirpath, _, filenames in os.walk(path):
for filename in filenames:
if os.path.splitext(filename)[-1].lower() in args.exts:
_path = os.path.join(dirpath, filename)
print(_path)
detect(sess, builder.model, builder.names, image, _path)
plt.show()
def main():
model = config.get('config', 'model')
yolo = importlib.import_module('model.' + model)
width = config.getint(model, 'width')
height = config.getint(model, 'height')
preprocess = getattr(importlib.import_module('detect'), args.preprocess)
with tf.Session() as sess:
ph_image = tf.placeholder(tf.float32, [1, height, width, 3], name='ph_image')
builder = yolo.Builder(args, config)
builder(ph_image)
global_step = tf.contrib.framework.get_or_create_global_step()
model_path = tf.train.latest_checkpoint(utils.get_logdir(config))
tf.logging.info('load ' + model_path)
slim.assign_from_checkpoint_fn(model_path, tf.global_variables())(sess)
tf.logging.info('global_step=%d' % sess.run(global_step))
tensors = [builder.model.conf, builder.model.xy_min, builder.model.xy_max]
tensors = [tf.check_numerics(t, t.op.name) for t in tensors]
cap = cv2.VideoCapture(0)
try:
while True:
ret, image_bgr = cap.read()
assert ret
image_height, image_width, _ = image_bgr.shape
scale = [image_width / builder.model.cell_width, image_height / builder.model.cell_height]
image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
image_std = np.expand_dims(preprocess(cv2.resize(image_rgb, (width, height))).astype(np.float32), 0)
feed_dict = {ph_image: image_std}
conf, xy_min, xy_max = sess.run(tensors, feed_dict)
boxes = utils.postprocess.non_max_suppress(conf[0], xy_min[0], xy_max[0], args.threshold, args.threshold_iou)
for _conf, _xy_min, _xy_max in boxes:
index = np.argmax(_conf)
if _conf[index] > args.threshold:
_xy_min = (_xy_min * scale).astype(np.int)
_xy_max = (_xy_max * scale).astype(np.int)
cv2.rectangle(image_bgr, tuple(_xy_min), tuple(_xy_max), (255, 0, 255), 3)
cv2.putText(image_bgr, builder.names[index] + ' (%.1f%%)' % (_conf[index] * 100), tuple(_xy_min), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
cv2.imshow('detection', image_bgr)
cv2.waitKey(1)
finally:
cv2.destroyAllWindows()
cap.release()
def __init__(self, content, style, content_names, style_names):
"""
Suppose the content and style is a numpy array,
"""
self.content_names = content_names
self.style_names = style_names
self.VGG_MEAN = [123.68, 116.78, 103.94]
tf.reset_default_graph()
content = tf.constant(content) - tf.reshape(tf.constant(self.VGG_MEAN), [1, 1, 3])
_, self.content_layers = nets.vgg.vgg_19(tf.expand_dims(content, axis = 0), is_training = False, spatial_squeeze = False)
layer_name, layer_value = zip(*filter(lambda x: x[0] in content_names, self.content_layers.items()))
init_fn = slim.assign_from_checkpoint_fn("./vgg_19.ckpt", slim.get_variables_to_restore())
with tf.Session() as s, tf.device("/device:XLA_CPU:0"):
init_fn(s)
layer_value = s.run(layer_value)
self.content_map = dict(zip(layer_name, layer_value))
#print(content_map)
tf.reset_default_graph()
style = tf.constant(style) - tf.reshape(tf.constant(self.VGG_MEAN), [1, 1, 3])
_, self.style_layers = nets.vgg.vgg_19(tf.expand_dims(style, axis = 0), is_training = False, spatial_squeeze = False)
layer_name, layer_value = zip(*filter(lambda x: x[0] in style_names, self.style_layers.items()))
init_fn = slim.assign_from_checkpoint_fn("./vgg_19.ckpt", slim.get_variables_to_restore())
with tf.Session() as s, tf.device("/device:XLA_CPU:0"):
init_fn(s)
layer_value = s.run(layer_value)
self.style_map = dict(zip(layer_name, layer_value))
#print(content_map)
tf.reset_default_graph()
self.target = tf.Variable(np.random.randint(0, 256, content.shape), dtype = tf.float32, name = "generate_image")
self._build_graph()
def train(self):
img_size = [self.image_height, self.image_width, self.image_depth]
train_batch = tf.train.shuffle_batch([read_tfrecord(self.train_file, img_size)],
batch_size = self.train_batch_size,
capacity = 3000,
num_threads = 2,
min_after_dequeue = 1000)
test_batch = tf.train.shuffle_batch([read_tfrecord(self.test_file, img_size)],
batch_size = self.test_batch_size,
capacity = 500,
num_threads = 2,
min_after_dequeue = 300)
init = tf.global_variables_initializer()
init_fn = slim.assign_from_checkpoint_fn("resnet_v2_50.ckpt", slim.get_model_variables('resnet_v2'))
saver = tf.train.Saver()
with tf.Session() as sess:
sess.run(init)
init_fn(sess)
train_writer = tf.summary.FileWriter(self.log_dir + "/train", sess.graph)
test_writer = tf.summary.FileWriter(self.log_dir + "/test", sess.graph)
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess=sess, coord=coord)
inputs_test, outputs_gt_test = build_img_pair(sess.run(test_batch))
for iter in range(self.max_iteration):
inputs_train, outputs_gt_train = build_img_pair(sess.run(train_batch))
# train with dynamic learning rate
if iter <= 500:
self.train_step.run({self.input_data:inputs_train, self.gt:outputs_gt_train,
self.learning_rate:1e-3, self.batch_size:self.train_batch_size})
elif iter <= self.max_iteration - 1000:
self.train_step.run({self.input_data:inputs_train, self.gt:outputs_gt_train,
self.learning_rate:0.5e-3, self.batch_size:self.train_batch_size})
else:
self.train_step.run({self.input_data:inputs_train, self.gt:outputs_gt_train,
self.learning_rate:1e-4, self.batch_size:self.train_batch_size})
# print training loss and test loss
if iter%10 == 0:
summary_train = sess.run(self.summary, {self.input_data:inputs_train, self.gt:outputs_gt_train,
self.batch_size:self.train_batch_size})
train_writer.add_summary(summary_train, iter)
train_writer.flush()
summary_test = sess.run(self.summary, {self.input_data:inputs_test, self.gt:outputs_gt_test,
self.batch_size:self.test_batch_size})
test_writer.add_summary(summary_test, iter)
test_writer.flush()
# record training loss and test loss
if iter%10 == 0:
train_loss = self.cross_entropy.eval({self.input_data:inputs_train, self.gt:outputs_gt_train,
self.batch_size:self.train_batch_size})
test_loss = self.cross_entropy.eval({self.input_data:inputs_test, self.gt:outputs_gt_test,
self.batch_size:self.test_batch_size})
print("iter step %d trainning batch loss %f"%(iter, train_loss))
print("iter step %d test loss %f\n"%(iter, test_loss))
# record model
if iter%100 == 0:
saver.save(sess, self.log_dir + "/model.ckpt", global_step=iter)
coord.request_stop()
coord.join(threads)
def train(self):
img_size = [self.image_height, self.image_width, self.image_depth]
train_batch = tf.train.shuffle_batch([read_tfrecord(self.train_file, img_size)],
batch_size = self.train_batch_size,
capacity = 3000,
num_threads = 2,
min_after_dequeue = 1000)
test_batch = tf.train.shuffle_batch([read_tfrecord(self.test_file, img_size)],
batch_size = self.test_batch_size,
capacity = 500,
num_threads = 2,
min_after_dequeue = 300)
init = tf.global_variables_initializer()
init_fn = slim.assign_from_checkpoint_fn("resnet_v2_50.ckpt", slim.get_model_variables('resnet_v2'))
saver = tf.train.Saver()
with tf.Session() as sess:
sess.run(init)
init_fn(sess)
train_writer = tf.summary.FileWriter(self.log_dir + "/train", sess.graph)
test_writer = tf.summary.FileWriter(self.log_dir + "/test", sess.graph)
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess=sess, coord=coord)
inputs_test, outputs_gt_test = build_img_pair(sess.run(test_batch))
for iter in range(self.max_iteration):
inputs_train, outputs_gt_train = build_img_pair(sess.run(train_batch))
# train with dynamic learning rate
if iter <= 500:
self.train_step.run({self.input_data:inputs_train, self.gt:outputs_gt_train,
self.learning_rate:1e-3, self.batch_size:self.train_batch_size})
elif iter <= self.max_iteration - 1000:
self.train_step.run({self.input_data:inputs_train, self.gt:outputs_gt_train,
self.learning_rate:0.5e-3, self.batch_size:self.train_batch_size})
else:
self.train_step.run({self.input_data:inputs_train, self.gt:outputs_gt_train,
self.learning_rate:1e-4, self.batch_size:self.train_batch_size})
# print training loss and test loss
if iter%10 == 0:
summary_train = sess.run(self.summary, {self.input_data:inputs_train, self.gt:outputs_gt_train,
self.batch_size:self.train_batch_size})
train_writer.add_summary(summary_train, iter)
train_writer.flush()
summary_test = sess.run(self.summary, {self.input_data:inputs_test, self.gt:outputs_gt_test,
self.batch_size:self.test_batch_size})
test_writer.add_summary(summary_test, iter)
test_writer.flush()
# record training loss and test loss
if iter%10 == 0:
train_loss = self.cross_entropy.eval({self.input_data:inputs_train, self.gt:outputs_gt_train,
self.batch_size:self.train_batch_size})
test_loss = self.cross_entropy.eval({self.input_data:inputs_test, self.gt:outputs_gt_test,
self.batch_size:self.test_batch_size})
print("iter step %d trainning batch loss %f"%(iter, train_loss))
print("iter step %d test loss %f\n"%(iter, test_loss))
# record model
if iter%100 == 0:
saver.save(sess, self.log_dir + "/model.ckpt", global_step=iter)
coord.request_stop()
coord.join(threads)
def train(self):
img_size = [self.image_height, self.image_width, self.image_depth]
train_batch = tf.train.shuffle_batch([read_tfrecord(self.train_file, img_size)],
batch_size = self.train_batch_size,
capacity = 2000,
num_threads = 2,
min_after_dequeue = 1000)
test_batch = tf.train.shuffle_batch([read_tfrecord(self.test_file, img_size)],
batch_size = self.test_batch_size,
capacity = 500,
num_threads = 2,
min_after_dequeue = 300)
init = tf.global_variables_initializer()
init_fn = slim.assign_from_checkpoint_fn("vgg_16.ckpt", slim.get_model_variables('vgg_16'))
saver = tf.train.Saver()
with tf.Session() as sess:
sess.run(init)
init_fn(sess)
train_writer = tf.summary.FileWriter(self.log_dir + "/train", sess.graph)
test_writer = tf.summary.FileWriter(self.log_dir + "/test", sess.graph)
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess=sess, coord=coord)
inputs_test, outputs_gt_test = build_img_pair(sess.run(test_batch))
for iter in range(self.max_iteration):
inputs_train, outputs_gt_train = build_img_pair(sess.run(train_batch))
# train with dynamic learning rate
if iter <= 500:
self.train_step.run({self.input_data:inputs_train, self.gt:outputs_gt_train, self.is_training:True,
self.learning_rate:1e-4})
elif iter <= self.max_iteration - 1000:
self.train_step.run({self.input_data:inputs_train, self.gt:outputs_gt_train, self.is_training:True,
self.learning_rate:0.5e-4})
else:
self.train_step.run({self.input_data:inputs_train, self.gt:outputs_gt_train, self.is_training:True,
self.learning_rate:1e-5})
# print training loss and test loss
if iter%10 == 0:
summary_train = sess.run(self.summary, {self.input_data:inputs_train, self.gt:outputs_gt_train,
self.is_training:False})
train_writer.add_summary(summary_train, iter)
train_writer.flush()
summary_test = sess.run(self.summary, {self.input_data:inputs_test, self.gt:outputs_gt_test,
self.is_training:False})
test_writer.add_summary(summary_test, iter)
test_writer.flush()
# record training loss and test loss
if iter%10 == 0:
train_loss = self.cross_entropy.eval({self.input_data:inputs_train, self.gt:outputs_gt_train,
self.is_training:False})
test_loss = self.cross_entropy.eval({self.input_data:inputs_test, self.gt:outputs_gt_test,
self.is_training:False})
print("iter step %d trainning batch loss %f"%(iter, train_loss))
print("iter step %d test loss %f\n"%(iter, test_loss))
# record model
if iter%100 == 0:
saver.save(sess, self.log_dir + "/model.ckpt", global_step=iter)
coord.request_stop()
coord.join(threads)
def train(self):
img_size = [self.image_height, self.image_width, self.image_depth]
train_batch = tf.train.shuffle_batch([read_tfrecord(self.train_file, img_size)],
batch_size = self.train_batch_size,
capacity = 3000,
num_threads = 2,
min_after_dequeue = 1000)
test_batch = tf.train.shuffle_batch([read_tfrecord(self.test_file, img_size)],
batch_size = self.test_batch_size,
capacity = 500,
num_threads = 2,
min_after_dequeue = 300)
init = tf.global_variables_initializer()
init_fn = slim.assign_from_checkpoint_fn("resnet_v2_50.ckpt", slim.get_model_variables('resnet_v2'))
saver = tf.train.Saver()
with tf.Session() as sess:
sess.run(init)
init_fn(sess)
train_writer = tf.summary.FileWriter(self.log_dir + "/train", sess.graph)
test_writer = tf.summary.FileWriter(self.log_dir + "/test", sess.graph)
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess=sess, coord=coord)
inputs_test, outputs_gt_test = build_img_pair(sess.run(test_batch))
for iter in range(self.max_iteration):
inputs_train, outputs_gt_train = build_img_pair(sess.run(train_batch))
# train with dynamic learning rate
if iter <= 500:
self.train_step.run({self.input_data:inputs_train, self.gt:outputs_gt_train,
self.learning_rate:1e-3, self.batch_size:self.train_batch_size})
elif iter <= self.max_iteration - 1000:
self.train_step.run({self.input_data:inputs_train, self.gt:outputs_gt_train,
self.learning_rate:0.5e-3, self.batch_size:self.train_batch_size})
else:
self.train_step.run({self.input_data:inputs_train, self.gt:outputs_gt_train,
self.learning_rate:1e-4, self.batch_size:self.train_batch_size})
# print training loss and test loss
if iter%10 == 0:
summary_train = sess.run(self.summary, {self.input_data:inputs_train, self.gt:outputs_gt_train,
self.batch_size:self.train_batch_size})
train_writer.add_summary(summary_train, iter)
train_writer.flush()
summary_test = sess.run(self.summary, {self.input_data:inputs_test, self.gt:outputs_gt_test,
self.batch_size:self.test_batch_size})
test_writer.add_summary(summary_test, iter)
test_writer.flush()
# record training loss and test loss
if iter%10 == 0:
train_loss = self.cross_entropy.eval({self.input_data:inputs_train, self.gt:outputs_gt_train,
self.batch_size:self.train_batch_size})
test_loss = self.cross_entropy.eval({self.input_data:inputs_test, self.gt:outputs_gt_test,
self.batch_size:self.test_batch_size})
print("iter step %d trainning batch loss %f"%(iter, train_loss))
print("iter step %d test loss %f\n"%(iter, test_loss))
# record model
if iter%100 == 0:
saver.save(sess, self.log_dir + "/model.ckpt", global_step=iter)
coord.request_stop()
coord.join(threads)
def train(self):
img_size = [self.image_height, self.image_width, self.image_depth]
train_batch = tf.train.shuffle_batch([read_tfrecord(self.train_file, img_size)],
batch_size = self.train_batch_size,
capacity = 2000,
num_threads = 2,
min_after_dequeue = 1000)
test_batch = tf.train.shuffle_batch([read_tfrecord(self.test_file, img_size)],
batch_size = self.test_batch_size,
capacity = 500,
num_threads = 2,
min_after_dequeue = 300)
init = tf.global_variables_initializer()
init_fn = slim.assign_from_checkpoint_fn("vgg_16.ckpt", slim.get_model_variables('vgg_16'))
saver = tf.train.Saver()
with tf.Session() as sess:
sess.run(init)
init_fn(sess)
train_writer = tf.summary.FileWriter(self.log_dir + "/train", sess.graph)
test_writer = tf.summary.FileWriter(self.log_dir + "/test", sess.graph)
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess=sess, coord=coord)
inputs_test, outputs_gt_test = build_img_pair(sess.run(test_batch))
for iter in range(self.max_iteration):
inputs_train, outputs_gt_train = build_img_pair(sess.run(train_batch))
# train with dynamic learning rate
if iter <= 500:
self.train_step.run({self.input_data:inputs_train, self.gt:outputs_gt_train,
self.learning_rate:1e-4, self.is_training:True})
elif iter <= self.max_iteration - 1000:
self.train_step.run({self.input_data:inputs_train, self.gt:outputs_gt_train,
self.learning_rate:0.5e-4, self.is_training:True})
else:
self.train_step.run({self.input_data:inputs_train, self.gt:outputs_gt_train,
self.learning_rate:1e-5, self.is_training:True})
# print training loss and test loss
if iter%10 == 0:
summary_train = sess.run(self.summary, {self.input_data:inputs_train, self.gt:outputs_gt_train, self.is_training:False})
train_writer.add_summary(summary_train, iter)
train_writer.flush()
summary_test = sess.run(self.summary, {self.input_data:inputs_test, self.gt:outputs_gt_test, self.is_training:False})
test_writer.add_summary(summary_test, iter)
test_writer.flush()
# record training loss and test loss
if iter%10 == 0:
train_loss = self.cross_entropy.eval({self.input_data:inputs_train, self.gt:outputs_gt_train, self.is_training:False})
test_loss = self.cross_entropy.eval({self.input_data:inputs_test, self.gt:outputs_gt_test, self.is_training:False})
print("iter step %d trainning batch loss %f"%(iter, train_loss))
print("iter step %d test loss %f\n"%(iter, test_loss))
# record model
if iter%100 == 0:
saver.save(sess, self.log_dir + "/model.ckpt", global_step=iter)
coord.request_stop()
coord.join(threads)
def train(self,
train_anchor_batch,
train_pos_batch,
train_neg_batch,
train_params,
preprocessed=True):
"""Training process of the matcher.
Each input data should have same shape.
Args:
train_anchor_batch: anchor batch.
train_pos_batch: positive batch.
train_neg_batch: negative batch.
train_params: commons.TrainTestParams object.
preprocessed: if data has been preprocessed.
"""
self.check_dm_model_exist()
self.dm_model.use_graph()
# get embedding for all batches.
all_batches = tf.concat(
0, [train_anchor_batch, train_pos_batch, train_neg_batch])
if not preprocessed:
all_batches = self.dm_model.preprocess(all_batches)
all_feats, _ = self.build_model(all_batches)
anchor_feats, pos_feats, neg_feats = tf.split(all_feats, 3, axis=0)
self.set_key_vars(train_params.restore_scopes_exclude,
train_params.train_scopes)
self.compute_losses(anchor_feats, pos_feats, neg_feats, train_params)
init_fn = None
if train_params.fine_tune:
# self.vars_to_restore is supposed to be set in set_key_vars
print("[dm_matcher.train: info] Trying to restore variables: {}".format(
self.vars_to_restore))
init_fn = slim.assign_from_checkpoint_fn(train_params.custom["model_fn"],
self.vars_to_restore)
if not train_params.resume_training:
data_manager.remove_dir(train_params.train_log_dir)
if train_params.use_regularization:
regularization_loss = tf.add_n(tf.losses.get_regularization_losses())
tf.summary.scalar("losses/regularization_loss", regularization_loss)
total_loss = tf.losses.get_total_loss(
add_regularization_losses=train_params.use_regularization)
base_model.train_model_given_loss(
total_loss, self.vars_to_train, train_params, init_fn=init_fn)
# TODO(jiefeng): to load weights from file.
def train(self,
train_input_batch,
train_label_batch,
train_params,
preprocessed=True):
"""Training process of the classifier.
Args:
train_input_batch: input batch for training.
train_label_batch: class id for training.
train_params: commons.TrainTestParams object.
preprocessed: if train data has been preprocessed.
"""
assert train_input_batch is not None, "train input batch is none"
assert train_label_batch is not None, "train label batch is none"
assert isinstance(
train_params,
commons.TrainTestParams), "train params is not a valid type"
self.check_dm_model_exist()
# self.dm_model.use_graph()
model_params = self.dm_model.net_params
if not preprocessed:
train_input_batch = self.dm_model.preprocess(train_input_batch)
pred_logits, endpoints = self.build_model(train_input_batch)
self.set_key_vars(train_params.restore_scopes_exclude,
train_params.train_scopes)
comp_train_accuracy(pred_logits, train_label_batch)
tf.assert_equal(
tf.reduce_max(train_label_batch),
tf.convert_to_tensor(
model_params.cls_num, dtype=tf.int64))
onehot_labels = tf.one_hot(
train_label_batch, model_params.cls_num, on_value=1.0, off_value=0.0)
# onehot_labels = slim.one_hot_encoding(train_label_batch,
# model_params.cls_num)
onehot_labels = tf.squeeze(onehot_labels)
self.compute_losses(onehot_labels, pred_logits, endpoints)
init_fn = None
if train_params.fine_tune and not train_params.resume_training:
init_fn = slim.assign_from_checkpoint_fn(train_params.custom["model_fn"],
self.vars_to_restore)
# this would not work if a tensorboard is running...
if not train_params.resume_training:
data_manager.remove_dir(train_params.train_log_dir)
# display regularization loss.
if train_params.use_regularization:
regularization_loss = tf.add_n(tf.losses.get_regularization_losses())
tf.summary.scalar("losses/regularization_loss", regularization_loss)
total_loss = tf.losses.get_total_loss(
add_regularization_losses=train_params.use_regularization)
base_model.train_model_given_loss(
total_loss, self.vars_to_train, train_params, init_fn=init_fn)
def __get_init_fn(self):
"""Returns a function run by the chief worker to warm-start the training.
Note that the init_fn is only run when initializing the model during the very
first global step.
Returns:
An init function run by the supervisor.
"""
if self.checkpoint_path is None:
return None
# Warn the user if a checkpoint exists in the train_dir. Then we'll be
# ignoring the checkpoint anyway.
if tf.train.latest_checkpoint(self.train_dir):
tf.logging.info(
'Ignoring --checkpoint_path because a checkpoint already exists in %s'
% self.train_dir)
return None
exclusions = []
if self.checkpoint_exclude_scopes:
exclusions = [scope.strip()
for scope in self.checkpoint_exclude_scopes.split(',')]
# TODO(sguada) variables.filter_variables()
variables_to_restore = []
for var in slim.get_model_variables():
excluded = False
for exclusion in exclusions:
if var.op.name.startswith(exclusion):
excluded = True
break
if not excluded:
variables_to_restore.append(var)
if tf.gfile.IsDirectory(self.checkpoint_path):
checkpoint_path = tf.train.latest_checkpoint(self.checkpoint_path)
else:
checkpoint_path = self.checkpoint_path
tf.logging.info('Fine-tuning from %s' % checkpoint_path)
return slim.assign_from_checkpoint_fn(
checkpoint_path,
variables_to_restore,
ignore_missing_vars=self.ignore_missing_vars)
def use_fined_model(self):
image_size = inception.inception_v4.default_image_size
batch_size = 3
flowers_data_dir = "../../data/flower"
train_dir = '/tmp/inception_finetuned/'
with tf.Graph().as_default():
tf.logging.set_verbosity(tf.logging.INFO)
dataset = flowers.get_split('train', flowers_data_dir)
images, images_raw, labels = self.load_batch(dataset, height=image_size, width=image_size)
# Create the model, use the default arg scope to configure the batch norm parameters.
with slim.arg_scope(inception.inception_v4_arg_scope()):
logits, _ = inception.inception_v4(images, num_classes=dataset.num_classes, is_training=True)
probabilities = tf.nn.softmax(logits)
checkpoint_path = tf.train.latest_checkpoint(train_dir)
init_fn = slim.assign_from_checkpoint_fn(
checkpoint_path,
slim.get_variables_to_restore())
with tf.Session() as sess:
with slim.queues.QueueRunners(sess):
sess.run(tf.initialize_local_variables())
init_fn(sess)
np_probabilities, np_images_raw, np_labels = sess.run([probabilities, images_raw, labels])
for i in range(batch_size):
image = np_images_raw[i, :, :, :]
true_label = np_labels[i]
predicted_label = np.argmax(np_probabilities[i, :])
predicted_name = dataset.labels_to_names[predicted_label]
true_name = dataset.labels_to_names[true_label]
plt.figure()
plt.imshow(image.astype(np.uint8))
plt.title('Ground Truth: [%s], Prediction [%s]' % (true_name, predicted_name))
plt.axis('off')
plt.show()
return
def __get_init_fn(self):
"""Returns a function run by the chief worker to warm-start the training.
Note that the init_fn is only run when initializing the model during the very
first global step.
Returns:
An init function run by the supervisor.
"""
if self.checkpoint_path is None:
return None
# Warn the user if a checkpoint exists in the train_dir. Then we'll be
# ignoring the checkpoint anyway.
if tf.train.latest_checkpoint(self.train_dir):
tf.logging.info(
'Ignoring --checkpoint_path because a checkpoint already exists in %s'
% self.train_dir)
return None
exclusions = []
if self.checkpoint_exclude_scopes:
exclusions = [scope.strip()
for scope in self.checkpoint_exclude_scopes.split(',')]
# TODO(sguada) variables.filter_variables()
variables_to_restore = []
all_variables = slim.get_model_variables()
if self.fine_tune_vgg16:
global_step = slim.get_or_create_global_step()
all_variables.append(global_step)
for var in all_variables:
excluded = False
for exclusion in exclusions:
if var.op.name.startswith(exclusion):
excluded = True
break
if not excluded:
variables_to_restore.append(var)
if tf.gfile.IsDirectory(self.checkpoint_path):
checkpoint_path = tf.train.latest_checkpoint(self.checkpoint_path)
else:
checkpoint_path = self.checkpoint_path
tf.logging.info('Fine-tuning from %s' % checkpoint_path)
return slim.assign_from_checkpoint_fn(
checkpoint_path,
variables_to_restore,
ignore_missing_vars=self.ignore_missing_vars)
def main():
model = config.get('config', 'model')
cachedir = utils.get_cachedir(config)
with open(os.path.join(cachedir, 'names'), 'r') as f:
names = [line.strip() for line in f]
width = config.getint(model, 'width')
height = config.getint(model, 'height')
yolo = importlib.import_module('model.' + model)
cell_width, cell_height = utils.calc_cell_width_height(config, width, height)
tf.logging.info('(width, height)=(%d, %d), (cell_width, cell_height)=(%d, %d)' % (width, height, cell_width, cell_height))
with tf.Session() as sess:
paths = [os.path.join(cachedir, profile + '.tfrecord') for profile in args.profile]
num_examples = sum(sum(1 for _ in tf.python_io.tf_record_iterator(path)) for path in paths)
tf.logging.warn('num_examples=%d' % num_examples)
image_rgb, labels = utils.data.load_image_labels(paths, len(names), width, height, cell_width, cell_height, config)
image_std = tf.image.per_image_standardization(image_rgb)
image_rgb = tf.cast(image_rgb, tf.uint8)
ph_image = tf.placeholder(image_std.dtype, [1] + image_std.get_shape().as_list(), name='ph_image')
global_step = tf.contrib.framework.get_or_create_global_step()
builder = yolo.Builder(args, config)
builder(ph_image)
variables_to_restore = slim.get_variables_to_restore()
ph_labels = [tf.placeholder(l.dtype, [1] + l.get_shape().as_list(), name='ph_' + l.op.name) for l in labels]
with tf.name_scope('total_loss') as name:
builder.create_objectives(ph_labels)
total_loss = tf.losses.get_total_loss(name=name)
tf.global_variables_initializer().run()
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess, coord)
_image_rgb, _image_std, _labels = sess.run([image_rgb, image_std, labels])
coord.request_stop()
coord.join(threads)
feed_dict = dict([(ph, np.expand_dims(d, 0)) for ph, d in zip(ph_labels, _labels)])
feed_dict[ph_image] = np.expand_dims(_image_std, 0)
logdir = utils.get_logdir(config)
assert os.path.exists(logdir)
model_path = tf.train.latest_checkpoint(logdir)
tf.logging.info('load ' + model_path)
slim.assign_from_checkpoint_fn(model_path, variables_to_restore)(sess)
tf.logging.info('global_step=%d' % sess.run(global_step))
tf.logging.info('total_loss=%f' % sess.run(total_loss, feed_dict))
_ = Drawer(sess, names, builder.model.cell_width, builder.model.cell_height, _image_rgb, _labels, builder.model, feed_dict)
plt.show()
def _get_init_fn():
"""Returns a function run by the chief worker to warm-start the training.
Note that the init_fn is only run when initializing the model during the very
first global step.
Returns:
An init function run by the supervisor.
"""
if FLAGS.checkpoint_path is None:
return None
# Warn the user if a checkpoint exists in the train_dir. Then we'll be
# ignoring the checkpoint anyway.
if tf.train.latest_checkpoint(FLAGS.train_dir):
tf.logging.info(
'Ignoring --checkpoint_path because a checkpoint already exists in %s'
% FLAGS.train_dir)
return None
exclusions = []
if FLAGS.checkpoint_exclude_scopes:
exclusions = [scope.strip()
for scope in FLAGS.checkpoint_exclude_scopes.split(',')]
# TODO(sguada) variables.filter_variables()
variables_to_restore = []
for var in slim.get_model_variables():
excluded = False
for exclusion in exclusions:
if var.op.name.startswith(exclusion):
excluded = True
break
if not excluded:
variables_to_restore.append(var)
if tf.gfile.IsDirectory(FLAGS.checkpoint_path):
checkpoint_path = tf.train.latest_checkpoint(FLAGS.checkpoint_path)
else:
checkpoint_path = FLAGS.checkpoint_path
tf.logging.info('Fine-tuning from %s' % checkpoint_path)
return slim.assign_from_checkpoint_fn(
checkpoint_path,
variables_to_restore,
ignore_missing_vars=FLAGS.ignore_missing_vars)
def export_to_h5(checkpoint_dir, export_path, images, end_points, num_samples,
batch_size, sact):
"""Exports ponder cost maps and other useful info to an HDF5 file."""
output_file = h5py.File(export_path, 'w')
output_file.attrs['block_scopes'] = end_points['block_scopes']
keys_to_tensors = {}
for block_scope in end_points['block_scopes']:
for k in ('{}/ponder_cost'.format(block_scope),
'{}/num_units'.format(block_scope),
'{}/halting_distribution'.format(block_scope),
'{}/flops'.format(block_scope)):
keys_to_tensors[k] = end_points[k]
keys_to_tensors['images'] = images
keys_to_tensors['flops'] = end_points['flops']
if sact:
keys_to_tensors['ponder_cost_map'] = sact_map(end_points, 'ponder_cost')
keys_to_tensors['num_units_map'] = sact_map(end_points, 'num_units')
keys_to_datasets = {}
for key, tensor in keys_to_tensors.iteritems():
sh = tensor.get_shape().as_list()
sh[0] = num_samples
print(key, sh)
keys_to_datasets[key] = output_file.create_dataset(
key, sh, compression='lzf')
variables_to_restore = slim.get_model_variables()
checkpoint_path = tf.train.latest_checkpoint(checkpoint_dir)
assert checkpoint_path is not None
init_fn = slim.assign_from_checkpoint_fn(checkpoint_path,
variables_to_restore)
sv = tf.train.Supervisor(
graph=tf.get_default_graph(),
logdir=None,
summary_op=None,
summary_writer=None,
global_step=None,
saver=None)
assert num_samples % batch_size == 0
num_batches = num_samples // batch_size
with sv.managed_session('', start_standard_services=False) as sess:
init_fn(sess)
sv.start_queue_runners(sess)
for i in range(num_batches):
tf.logging.info('Evaluating batch %d/%d', i + 1, num_batches)
end_points_out = sess.run(keys_to_tensors)
for key, dataset in keys_to_datasets.iteritems():
dataset[i * batch_size:(i + 1) * batch_size, ...] = end_points_out[key]
