def load_resnet_v2_152(model_dir, sess):
model_url = "http://download.tensorflow.org/models/resnet_v2_152_2017_04_14.tar.gz"
filename = model_url.split("/")[-1]
filepath = os.path.join(model_dir, filename.split(".tgz")[0])
try:
utils.download_pretrained_model_weights(model_url,
filepath,
unzip=True)
except:
print("Pre-training weights download failed!")
model_file_name = "resnet_v2_152.ckpt"
model_path = os.path.join(filepath, model_file_name)
resized_input_tensor = tf.placeholder(tf.float32,
shape=[None, None, None, 3])
with tf.contrib.slim.arg_scope(resnet_v2.resnet_arg_scope()):
bottleneck_tensor, _ = resnet_v2.resnet_v2_152(resized_input_tensor,
num_classes=None,
global_pool=True)
variable_restore_op = tf.contrib.slim.assign_from_checkpoint_fn(
model_path,
tf.contrib.slim.get_trainable_variables(),
ignore_missing_vars=True)
variable_restore_op(sess)
bottleneck_tensor = tf.squeeze(bottleneck_tensor, axis=[1, 2])
bottleneck_tensor_size = 2048
return bottleneck_tensor, resized_input_tensor, bottleneck_tensor_size
def load_resnet_v2_50(model_dir, sess):
model_file_name = "resnet_v2_50_2017_04_14/resnet_v2_50.ckpt"
model_path = os.path.join(model_dir, model_file_name)
resized_input_tensor = tf.placeholder(tf.float32, shape=[None, None, None, 3])
with tf.contrib.slim.arg_scope(resnet_v2.resnet_arg_scope()):
bottleneck_tensor, _ = resnet_v2.resnet_v2_50(
resized_input_tensor, num_classes=None, global_pool=True)
variable_restore_op = tf.contrib.slim.assign_from_checkpoint_fn(
model_path,
tf.contrib.slim.get_trainable_variables(),
ignore_missing_vars=True)
variable_restore_op(sess)
bottleneck_tensor = tf.squeeze(bottleneck_tensor, axis=[1, 2])
bottleneck_tensor_size = 2048
return bottleneck_tensor, resized_input_tensor, bottleneck_tensor_size