def restore(sess, global_vars):
print('from vgg_16 pretrained model')
reader = NewCheckpointReader(os.path.join(os.getcwd(),
'model/vgg_16.ckpt'))
# no batchnorm from vgg_16 pretrained model
restored_var_names = [
name + ':0' for name in reader.get_variable_to_dtype_map().keys()
if re.match('^.*weights$', name)
] # skip conv's biases
restored_vars = [
var for var in global_vars if var.name in restored_var_names
]
restored_var_names = [var.name[:-2] for var in restored_vars]
value_ph = tf.placeholder(dtype=tf.float32)
for i in range(len(restored_var_names)):
sess.run(
tf.assign(restored_vars[i], value_ph),
feed_dict={value_ph: reader.get_tensor(restored_var_names[i])})
initialized_vars = [var for var in global_vars if not var in restored_vars]
sess.run(tf.variables_initializer(initialized_vars))
def restore(sess, global_vars):
print('from resnet_v2_50 pretrained model')
reader = NewCheckpointReader(os.path.join(
os.getcwd(), 'model/resnet_v2_50.ckpt'))
# restore both weights and biases from conv and shortcut layers
restored_var_names = [name + ':0'
for name in reader.get_variable_to_dtype_map().keys()
if re.match('^.*weights$', name) or re.match('^.*biases$', name)]
restored_vars = [var for var in global_vars
if var.name in restored_var_names]
restored_var_names = [var.name[:-2] for var in restored_vars]
value_ph = tf.placeholder(dtype=tf.float32)
for i in range(len(restored_var_names)):
sess.run(tf.assign(restored_vars[i], value_ph),
feed_dict={value_ph: reader.get_tensor(restored_var_names[i])})
initialized_vars = [var for var in global_vars
if not var in restored_vars]
sess.run(tf.variables_initializer(initialized_vars))
def load_ckpt(self, pretrained):
# restore model with ckpt/pretrain or init
try:
print('trying to restore last checkpoint')
last_ckpt_path = tf.train.latest_checkpoint(
checkpoint_dir=cfg.ckpt_dir)
self.saver.restore(self.sess, save_path=last_ckpt_path)
print('restored checkpoint from:', last_ckpt_path)
except:
if self.is_training:
print('init variables')
restored_vars = []
global_vars = tf.global_variables()
if pretrained: # restore from tf-slim model
if os.path.exists(
os.path.join(cfg.workspace, premodel['path'])):
print('from ' + premodel['endp'])
import re
from tensorflow.python.pywrap_tensorflow import NewCheckpointReader
reader = NewCheckpointReader(
os.path.join(cfg.workspace, premodel['ckpt']))
# only restoring conv's weights
restored_var_names = [
name + ':0' for name in
reader.get_variable_to_dtype_map().keys()
if re.match(premodel['rptn'], name)
]
# update restored variables from pretrained model
restored_vars = [
var for var in global_vars
if var.name in restored_var_names
]
# update restored variables' name
restored_var_names = [
var.name[:-2] for var in restored_vars
]
# assignment variables
value_ph = tf.placeholder(tf.float32, shape=None)
for i in range(len(restored_var_names)):
self.sess.run(
tf.assign(restored_vars[i], value_ph),
feed_dict={
value_ph:
reader.get_tensor(restored_var_names[i])
})
initialized_vars = list(set(global_vars) - set(restored_vars))
self.sess.run(tf.variables_initializer(initialized_vars))
def get_list_of_variables_from_ckpt(ckpt_file):
reader = NewCheckpointReader(ckpt_file)
names = reader.get_variable_to_dtype_map()
return list(names.keys())