def pytorch2savedmodel(onnx_model_path, saved_model_dir):
onnx_model = onnx.load(onnx_model_path)
input_names = ["image_array"]
k_model = onnx_to_keras(
onnx_model=onnx_model,
input_names=input_names,
change_ordering=True,
verbose=False,
)
weights = k_model.get_weights()
K.set_learning_phase(0)
saved_model_dir = Path(saved_model_dir)
if saved_model_dir.exists():
shutil.rmtree(str(saved_model_dir))
saved_model_dir.mkdir()
tf.saved_model.save(
k_model,
str(saved_model_dir.joinpath("1")),
)
def __call__(self, inputs, conditioning_inputs, mode=tf.estimator.ModeKeys.TRAIN):
set_learning_phase(True) if mode == tf.estimator.ModeKeys.TRAIN else set_learning_phase(False)
self._build_layers()
masks = inputs
images = conditioning_inputs['images']
concat_images_masks = tf.concat([images, masks], axis=self.channel_axis)
# pass the inputs through the rib blocks
left1, right1, center1 = self.rib1(images, masks, concat_images_masks)
left2, right2, center2 = self.rib2(left1, right1, center1)
left3, right3, center3 = self.rib3(left2, right2, center2)
# flatten the output from the rib blocks using global average pooling
ribs_output = tf.concat([left3, right3, center3], axis=self.channel_axis)
ribs_output = self.conv_to_fc(ribs_output)
# pass the results through to the dense layers
ribs_output = self.dense1(ribs_output)
ribs_output = self.relu(ribs_output)
ribs_output = self.dense2(ribs_output)
ribs_output = self.relu(ribs_output)
ribs_output = self.dense3(ribs_output)
return ribs_output
def main():
config = Config()
K.set_learning_phase(0) # make sure its testing mode
sess = K.get_session()
# Load keras model to find the output node name
# model = load_fsanet_model(config=config)
# print(config.Age_h5_checkpoint_name)
print("[INFO] Load model from </ {} />".format(config.gender_h5_path))
# Gender
model = load_model(config.gender_h5_path, custom_objects={'swish': swish}, compile=False)
# Emotion
# model = load_model(config.expr_h5_path, custom_objects={'swish': swish}, compile=False)
# Age estimation model
# model = load_model(config.age_h5_checkpoint, custom_objects={'swish': swish, 'tf': tf}, compile=False)
# Action model
# model = load_model(args.checkpoint_dir,
# custom_objects={'swish': swish, 'GlorotUniform': glorot_uniform, 'f1': f1, 'tf': tf},
# compile=False)
converted_output_node_names = [node.op.name for node in model.outputs]
print("[INFO] in: ", model.inputs)
print("[INFO] out: ", model.outputs)
constant_graph = tf.compat.v1.graph_util.convert_variables_to_constants(
sess,
sess.graph.as_graph_def(),
converted_output_node_names)
if not os.path.isdir(config.freeze_folder):
os.mkdir("config.age_freeze_folder")
name = os.path.split(config.gender_pb_path)[-1].split(".")[0]
tf.io.write_graph(constant_graph, config.freeze_folder, '%s.pbtxt' % name, as_text=True)
tf.io.write_graph(constant_graph, config.freeze_folder, '%s.pb' % name, as_text=False)
print(colored("[INFO] convert model is success, saved at: ", "cyan", attrs=['bold']))
print(config.freeze_folder + '/%s.pb' % name)
def create_pb_model_of_pretrained(Net):
K.set_learning_phase(0)
#if not(Net=='VGG'):
tf.reset_default_graph()
with K.get_session().as_default():
if Net == 'InceptionV1_slim':
model = InceptionV1_slim(include_top=True, weights='imagenet')
name = "tf_inception_v1_slim.pb"
elif Net == 'InceptionV1':
model = inception_v1_oldTF(weights='imagenet', include_top=True)
name = "tf_inception_v1.pb"
elif Net == 'VGG':
model = tf.keras.applications.vgg19.VGG19(include_top=False, weights='imagenet',\
input_shape=(224,224,3))
name = "tf_vgg19.pb"
elif Net == 'ResNet50':
model = tf.keras.applications.resnet50.ResNet50(include_top=True, weights='imagenet',\
input_shape= (224, 224, 3))
name = "tf_resnet50.pb"
else:
raise (ValueError(Net + ' is unknown'))
os.makedirs('./model', exist_ok=True)
frozen_graph = freeze_session(
K.get_session(),
output_names=[out.op.name for out in model.outputs],
clear_devices=True)
# Save the pb model
tf.io.write_graph(frozen_graph,
logdir="model",
name=name,
as_text=False)
def build_teacher_model(self, rgb, num_classes, k, n):
K.set_learning_phase(True)
nStages = [16, 16 * k, 32 * k, 64 * k]
conv1 = self.Convolution_without_bias(rgb,
self.num_channels,
nStages[0],
stride=1,
padding=1)
#print(conv1)
group1 = self.group(conv1, nStages[0], nStages[1], n, 1)
group2 = self.group(group1, nStages[1], nStages[2], n, 2)
group3 = self.group(group2, nStages[2], nStages[3], n, 2)
batchNorm = BatchNormalization(axis=-1,
name='BatchNormal',
trainable=self.trainable)(group3)
relu = tf.nn.relu(batchNorm, name='relu')
averagePool = tf.nn.avg_pool(relu,
ksize=[1, 8, 8, 1],
strides=[1, 1, 1, 1],
padding='VALID',
name='averagePool')
#print(averagePool)
self.fc = self.FullyConnect(averagePool, num_classes)
#print(self.fc)
self.softmax = tf.nn.softmax(self.fc)
#print(self.softmax)
return self
def onnx2keras_pb(onnx_model_path,
input_names,
output_names,
output_path,
swap_channel_ordering=True):
from onnx2keras import onnx_to_keras
from convert.utils import freeze_session
output_dir, filename = os.path.split(output_path)
onnx_model = onnx.load(onnx_model_path)
k_model = onnx_to_keras(onnx_model,
input_names,
change_ordering=swap_channel_ordering,
name_policy='renumerate')
weights = k_model.get_weights()
K.set_learning_phase(0)
with K.get_session() as sess:
init = tf.global_variables_initializer()
sess.run(init)
k_model.set_weights(weights)
frozen_graph = freeze_session(sess,
keep_var_names=None,
output_names=output_names)
tf.train.write_graph(frozen_graph, output_dir, filename, as_text=False)
def test_dbn2():
# tf.enable_eager_execution()
inputs = tf.keras.Input((7, 7, 16))
data = np.random.normal(0, 1, [1, 7, 7, 16])
data2 = np.random.normal(0, 1, [256, 7, 7, 16])
K.set_learning_phase(1)
# tf.set_random_seed(1)
decor1 = DecorelationNormalization(m_per_group=2,
decomposition='zca',
instance_norm=0)
decor2 = DecorelationNormalization(m_per_group=2,
decomposition='zca',
instance_norm=1)
out1, out2 = decor1(inputs), decor2(inputs)
op1 = K.function(inputs, out1)
op2 = K.function(inputs, out2)
import time
out2 = op2(data2)
t1 = time.time()
out1 = op1(data)
t2 = time.time()
out2 = op2(data)
t3 = time.time()
distance = np.sum(np.square(out1 - out2))
print(distance)
print('t1:', t2 - t1)
print('t2:', t3 - t2)
def _clone_and_build_model(mode,
keras_model,
custom_objects,
features=None,
labels=None):
"""Clone and build the given keras_model.
Args:
mode: training mode.
keras_model: an instance of compiled keras model.
custom_objects: Dictionary for custom objects.
features: Dict of tensors.
labels: Dict of tensors, or single tensor instance.
Returns:
The newly built model.
"""
# Set to True during training, False for inference or testing.
K.set_learning_phase(mode == model_fn_lib.ModeKeys.TRAIN)
input_tensors, target_tensors = _convert_estimator_io_to_keras(
keras_model, features, labels)
return models.clone_and_build_model(
keras_model,
input_tensors,
target_tensors,
custom_objects,
compile_clone=(mode != model_fn_lib.ModeKeys.PREDICT),
in_place_reset=(not keras_model._is_graph_network))
def convert_ckpt2pb_model(ckpt_path, pb_path):
"""Convert a *.ckpt model to *.pb model.
Args:
* ckpt_path: *.ckpt model's file path
* pb_path: *.pb model's file path
Returns:
* pb_info: *.pb model's input/output information
"""
print(tf.__version__, tf.__path__)
with tf.Graph().as_default() as graph:
# model definition
K.set_learning_phase(0)
if FLAGS.enbl_dst:
with tf.variable_scope('distilled_model'):
model_dst = MicroscopeModel(data_format=FLAGS.data_format)
with tf.variable_scope(FLAGS.model_scope):
model = MicroscopeModel(data_format=FLAGS.data_format)
vars_all = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES,
scope=FLAGS.model_scope)
print("vars_all", vars_all)
# display input/output tensors' names
tf.logging.info('inputs: {}'.format([x.name for x in model.inputs]))
tf.logging.info('outputs: {}'.format([x.name for x in model.outputs]))
assert len(model.inputs) == 1 and len(model.outputs) == 1
pb_info = {
'input_name': model.inputs[0].name,
'output_name': model.outputs[0].name,
}
# # restore model weights from the *.ckpt model
# sess = create_session()
# saver = tf.train.Saver(vars_all)
# graph = tf.get_default_graph()
# # print("==============",graph.get_tensor_by_name("batch_normalization/beta:0"))
# # ckpt_path = "/home/terse/code/python/PocketFlow/models/model_pruned/model.ckpt-10000.meta"
# print("================",ckpt_path,vars_all)
# saver.restore(sess, ckpt_path)
# tf.logging.info('model weights restored from: ' + ckpt_path)
# # apply fake pruning
# if FLAGS.enbl_fake_prune:
# apply_fake_pruning(sess)
# # write the graph to a *.pb file
# graph_def = graph.as_graph_def()
# graph_def = tf.graph_util.convert_variables_to_constants(
# sess, graph_def, [pb_info['output_name'].replace(':0', '')])
# tf.train.write_graph(
# graph_def, os.path.dirname(pb_path), os.path.basename(pb_path), as_text=False)
# tf.logging.info('*.pb model generated: ' + pb_path)
# # test the *.pb model
# test_pb_model(pb_path, pb_info)
return pb_info
def onnx2keras(onnx_model_path,
input_names,
output_dir,
swap_channel_ordering=True):
from onnx2keras import onnx_to_keras
onnx_model = onnx.load(onnx_model_path)
k_model = onnx_to_keras(onnx_model,
input_names,
change_ordering=swap_channel_ordering,
name_policy='renumerate')
weights = k_model.get_weights()
K.set_learning_phase(0)
with K.get_session() as sess:
init = tf.global_variables_initializer()
sess.run(init)
k_model.set_weights(weights)
tf.saved_model.simple_save(
sess,
output_dir,
inputs={
input.name: tensor
for input, tensor in zip(onnx_model.graph.input,
k_model.inputs)
},
outputs={
output.name: tensor
for output, tensor in zip(onnx_model.graph.output,
k_model.outputs)
})
def pytorch2savedmodel(onnx_model_path, saved_model_dir):
onnx_model = onnx.load(onnx_model_path)
input_names = ['input_ids', "attention_mask", "token_type_ids", "task_id"]
k_model = onnx_to_keras(onnx_model=onnx_model,
input_names=input_names,
change_ordering=True,
verbose=False)
weights = k_model.get_weights()
K.set_learning_phase(0)
saved_model_dir = Path(saved_model_dir)
if saved_model_dir.exists():
shutil.rmtree(str(saved_model_dir))
saved_model_dir.mkdir()
with K.get_session() as sess:
init = tf.global_variables_initializer()
sess.run(init)
k_model.set_weights(weights)
tf.saved_model.simple_save(
sess,
str(saved_model_dir.joinpath('1')),
inputs={'image_array': k_model.input},
outputs=dict((output.name, tensor) for output, tensor in zip(
onnx_model.graph.output, k_model.outputs)))
def build_student_conv2fc1(self, images, num_classes, temp_softmax):
K.set_learning_phase(True)
num_filters = [64, 128, 256, 512, 512]
with tf.name_scope('mentee'):
self.conv1_1 = self.build_student_oneConvLayer(
images, "conv1_1", num_filters[0])
pool1 = tf.nn.max_pool(self.conv1_1,
ksize=[1, 2, 2, 1],
strides=[1, 2, 2, 1],
padding='SAME',
name='pool1')
print(pool1)
self.conv2_1 = self.build_student_oneConvLayer(
pool1, "conv2_1", num_filters[1])
pool2 = tf.nn.max_pool(self.conv2_1,
ksize=[1, 2, 2, 1],
strides=[1, 2, 2, 1],
padding='SAME',
name='pool2')
print(pool2)
self.fc3 = self.fc_student(pool2, "fc3", num_classes)
self.softmax = tf.nn.softmax(self.fc3 / temp_softmax)
print(self.fc3)
return self
def build_resnet_conv1Block1Fc1(self, rgb, num_classes, k):
print("build_resnet_conv1Block1Fc1")
K.set_learning_phase(True)
nStages = [16, 16 * k, 32 * k, 64 * k]
conv1 = self.Convolution(rgb,
self.num_channels,
nStages[0],
stride=1,
padding=1)
print(conv1)
block = self.basic_block(conv1, nStages[0], nStages[1], 2)
batchNorm = BatchNormalization(axis=-1, name='BatchNormal')(block)
relu = tf.nn.relu(batchNorm, name='relu')
print(relu)
averagePool = tf.nn.avg_pool(relu,
ksize=[1, 8, 8, 1],
strides=[1, 1, 1, 1],
padding='SAME',
name='averagePool')
print(averagePool)
self.fc = self.FullyConnect(averagePool, num_classes)
print(self.fc)
self.softmax = tf.nn.softmax(self.fc)
print(self.softmax)
return self
def build_student_conv5fc1(self, images, num_classes, temp_softmax):
print("build_student_conv5fc1")
K.set_learning_phase(True)
#num_filters = [64, 128, 256, 512, 512] # origin
#num_filters = [64, 128, 256, 512, 512-96] # 100per, 90per
#num_filters = [64-2, 128-6, 256-10, 512-30, 512-258] # 70per
#num_filters = [64-6, 128-26, 256-46, 512-132, 512-258] # 60per
num_filters = [64-22, 128-68, 256-136, 512-306, 512-416] # 50per
#num_filters = [64-58, 128-122, 256-246, 512-486, 512-416] # 30per
with tf.name_scope('mentee'):
self.conv1_1 = self.build_student_oneConvLayer(images, "conv1_1", num_filters[0])
pool1 = tf.nn.max_pool(self.conv1_1, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME', name='pool1')
print(pool1)
self.conv2_1 = self.build_student_oneConvLayer(pool1, "conv2_1", num_filters[1])
pool2 = tf.nn.max_pool(self.conv2_1, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME', name='pool2')
print(pool2)
self.conv3_1 = self.build_student_oneConvLayer(pool2, "conv3_1", num_filters[2])
pool3 = tf.nn.max_pool(self.conv3_1, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME', name='pool3')
print(pool3)
self.conv4_1 = self.build_student_oneConvLayer(pool3, "conv4_1", num_filters[3])
pool4 = tf.nn.max_pool(self.conv4_1, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME', name='pool4')
print(pool4)
self.conv5_1 = self.build_student_oneConvLayer(pool4, "conv5_1", num_filters[4])
pool5 = tf.nn.max_pool(self.conv5_1, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME', name='pool5')
print(pool5)
self.fc3 = self.fc_student(pool5, "fc3", num_classes)
self.softmax = tf.nn.softmax(self.fc3 / temp_softmax)
print(self.fc3)
return self
def export_model(keras_model, export_path, model_version=0, weights_path=None):
"""Export a model for use with tensorflow-serving.
Args:
keras_model: instantiated Keras model to export
export_path: destination to save the exported model files
model_version: integer version of the model
weights_path: path to a .h5 or .tf weights file for the model to load
"""
# Start the tensorflow session
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.8, allow_growth=False)
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))
K.set_session(sess)
K._LEARNING_PHASE = tf.constant(0)
K.set_learning_phase(0)
# Create export path if it doesn't exist
export_path = os.path.join(export_path, str(model_version))
builder = SavedModelBuilder(export_path)
# legacy_init_op = tf.group(tf.tables_initializer(), name='legacy_init_op')
# Initialize global variables and the model
init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer())
sess.run(init_op)
# Load the model and the weights
if weights_path is not None:
keras_model.load_weights(weights_path)
if type(keras_model.input) is list:
output = keras_model.output[-1]
else:
output = keras_model.output
# Define prediction signature
if type(keras_model.input) is list:
input_map = {"input{}".format(i): input_tensor
for i, input_tensor in enumerate(keras_model.input)}
output_map = {"prediction": output}
else:
input_map = {"input": keras_model.input}
output_map = {"prediction": output}
prediction_signature = tf.saved_model.signature_def_utils.predict_signature_def(
input_map,
output_map
)
# Add the meta_graph and the variables to the builder
builder.add_meta_graph_and_variables(
sess, [tag_constants.SERVING],
signature_def_map={
signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
prediction_signature
})
# Save the graph
builder.save()
def cnn_model_fn(features, labels, mode):
"""Model function for CNN."""
# Input Layer
# Reshape X to 4-D tensor: [batch_size, width, height, channels]
# MNIST images are 28x28 pixels, and have one color channel
input_layer = tf.reshape(features["pixels"], [-1, 28, 28, 1])
if mode == tf.estimator.ModeKeys.TRAIN:
K.set_learning_phase(1)
else:
K.set_learning_phase(0)
conv1 = Convolution2D(32, (5, 5), activation='relu', input_shape=(28,28,1))(input_layer)
pool1 = MaxPooling2D(pool_size=(2,2))(conv1)
conv2 = Convolution2D(64, (5, 5), activation='relu')(pool1)
pool2 = MaxPooling2D(pool_size=(2,2))(conv2)
pool2_flat = Flatten()(pool2)
dense = Dense(1024, activation='relu')(pool2_flat)
dropout = Dropout(0.4)(dense)
logits = Dense(10, activation='linear')(dropout)
predictions = {
# Generate predictions (for PREDICT and EVAL mode)
"classes": tf.argmax(input=logits, axis=1),
# Add `softmax_tensor` to the graph. It is used for PREDICT and by the
# `logging_hook`.
"probabilities": tf.nn.softmax(logits, name="softmax_tensor")
}
prediction_output = tf.estimator.export.PredictOutput({"classes": tf.argmax(input=logits, axis=1),
"probabilities": tf.nn.softmax(logits, name="softmax_tensor")})
if mode == tf.estimator.ModeKeys.PREDICT:
return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions,
export_outputs={tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: prediction_output})
# Calculate Loss (for both TRAIN and EVAL modes)
onehot_labels = tf.one_hot(indices=tf.cast(labels, tf.int32), depth=10)
loss = tf.losses.softmax_cross_entropy(
onehot_labels=onehot_labels, logits=logits)
tf.summary.scalar('loss', loss)
tf.summary.histogram('conv1', conv1)
tf.summary.histogram('dense', dense)
# Configure the Training Op (for TRAIN mode)
if mode == tf.estimator.ModeKeys.TRAIN:
optimizer = tf.train.AdamOptimizer(learning_rate=1e-4)
train_op = optimizer.minimize(
loss=loss,
global_step=tf.train.get_global_step())
return tf.estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op)
# Add evaluation metrics (for EVAL mode)
eval_metric_ops = {
"accuracy": tf.metrics.accuracy(
labels=labels, predictions=predictions["classes"])}
return tf.estimator.EstimatorSpec(
mode=mode, loss=loss, eval_metric_ops=eval_metric_ops)
def build_student_conv5fc2(self, images, num_classes, temp_softmax):
print("build_student_conv5fc2")
K.set_learning_phase(True)
num_filters = [64, 128, 256, 512, 512]
with tf.name_scope('mentee'):
self.conv1_1 = self.build_student_oneConvLayer(
images, "conv1_1", num_filters[0])
pool1 = tf.nn.max_pool(self.conv1_1,
ksize=[1, 2, 2, 1],
strides=[1, 2, 2, 1],
padding='SAME',
name='pool1')
print(pool1)
self.conv2_1 = self.build_student_oneConvLayer(
pool1, "conv2_1", num_filters[1])
pool2 = tf.nn.max_pool(self.conv2_1,
ksize=[1, 2, 2, 1],
strides=[1, 2, 2, 1],
padding='SAME',
name='pool2')
print(pool2)
self.conv3_1 = self.build_student_oneConvLayer(
pool2, "conv3_1", num_filters[2])
pool3 = tf.nn.max_pool(self.conv3_1,
ksize=[1, 2, 2, 1],
strides=[1, 2, 2, 1],
padding='SAME',
name='pool3')
print(pool3)
self.conv4_1 = self.build_student_oneConvLayer(
pool3, "conv4_1", num_filters[3])
pool4 = tf.nn.max_pool(self.conv4_1,
ksize=[1, 2, 2, 1],
strides=[1, 2, 2, 1],
padding='SAME',
name='pool4')
print(pool4)
self.conv5_1 = self.build_student_oneConvLayer(
pool4, "conv5_1", num_filters[4])
pool5 = tf.nn.max_pool(self.conv5_1,
ksize=[1, 2, 2, 1],
strides=[1, 2, 2, 1],
padding='SAME',
name='pool5')
print(pool5)
fc1 = self.fc_student(pool5, "fc1", 4096)
self.fc3 = self.fc_student(fc1, "fc3", num_classes)
self.softmax = tf.nn.softmax(self.fc3 / temp_softmax)
print(self.fc3)
return self
def build_conv1fc1(self, rgb, num_classes):
print("build_conv1fc1")
K.set_learning_phase(True)
conv = self.Convolution(rgb,
self.num_channels,
16,
stride=1,
padding=1)
relu = tf.nn.relu(conv, name='relu')
self.fc = self.FullyConnect(relu, num_classes)
self.softmax = tf.nn.softmax(self.fc)
return self
def test_dbn():
inputs = tf.keras.layers.Input([8, 8, 16])
data = np.random.normal(0, 1, [128, 8, 8, 16])
# K.set_learning_phase(1)
decor = DecorelationNormalization(group=1, instance_norm=1)
out = decor(inputs)
out = tf.reduce_mean(out)
sess = K.get_session()
sess.run(tf.global_variables_initializer())
K.set_learning_phase(1)
outputs = sess.run([out], feed_dict={inputs: data})
print(np.mean(outputs))
def _clone_and_build_model(mode,
keras_model,
custom_objects,
features=None,
labels=None,
optimizer_config=None):
"""Clone and build the given keras_model.
Args:
mode: training mode.
keras_model: an instance of compiled keras model.
custom_objects: Dictionary for custom objects.
features: Dict of tensors.
labels: Dict of tensors, or single tensor instance.
optimizer_config: Optimizer config dictionary, returned by
`optimizer.get_config()`. This is used when cloning a model with
an optimizer. Since `_clone_and_build_model` is called in a different
graph and session from the model, `optimizer.get_config()` may raise an
error during the attempt to serialize the optimizer hyperparameter values.
Returns:
The newly built model.
"""
# Set to True during training, False for inference or testing.
K.set_learning_phase(mode == ModeKeys.TRAIN)
input_tensors, target_tensors, sample_weight_tensors = (
_convert_estimator_io_to_keras(keras_model, features, labels))
compile_clone = (mode != ModeKeys.PREDICT)
global_step = None
if compile_clone:
# Set iterations to the global step created by tf.train.create_global_step()
# which is automatically run in the estimator framework.
global_step = tf.compat.v1.train.get_or_create_global_step()
K.track_variable(global_step)
clone = models.clone_and_build_model(
keras_model,
input_tensors,
target_tensors,
custom_objects,
compile_clone=compile_clone,
in_place_reset=(not keras_model._is_graph_network),
optimizer_iterations=global_step,
optimizer_config=optimizer_config)
if sample_weight_tensors is not None:
sample_weight_tensors = training_utils.standardize_sample_weights(
sample_weight_tensors, clone.output_names)
# Update calculated loss (model.total_loss) to include sample weights.
clone._compile_weights_loss_and_weighted_metrics(sample_weight_tensors)
return clone
def __call__(self, features, mode=tf.estimator.ModeKeys.TRAIN):
set_learning_phase(True) if mode == tf.estimator.ModeKeys.TRAIN else set_learning_phase(False)
self._build_layers()
images = features['images']
x = self.conv1(images)
x = self.conv2(x)
x = self.conv3(x)
x = self.conv4(x)
x = self.conv5(x)
return x
def __init__(self, model, sess=None):
K.set_learning_phase(0)
self.model = model
with self.model.graph.as_default():
with self.model.session.as_default():
self.class_grads = K.function([self.model.model.input],
K.gradients(
self.model.model.output,
self.model.model.input))
self.out = K.function([self.model.model.input],
self.model.model.output)
def freeze_keras_model_to_pb(tk_model,
export_dir='.',
export_name=None,
optimize_graph=True,
verbose=True):
"""
not safe
Known issue:
run into 'SystemError: unknown opcode' if your keras model contains lambda layers
:param tk_model:
:param export_dir:
:param export_name:
:param optimize_graph:
:param verbose:
:return:
"""
sess_or = K.get_session()
weights = tk_model.get_weights()
try:
with tf.Graph().as_default() as graph, tf.Session(
graph=graph).as_default() as sess:
K.set_session(sess)
K.set_learning_phase(0)
new_model = tf.keras.models.clone_model(tk_model)
new_model.trainable = False
sess.run(tf.global_variables_initializer())
new_model.set_weights(weights)
input_names = [item.name for item in new_model.inputs]
output_names = [item.name for item in new_model.outputs]
graph = freeze_sess_to_constant_pb(sess,
output_node_names=output_names)
if export_name:
export_name = export_name.strip('.pb')
else:
export_name = tk_model.name
tf.io.write_graph(graph,
export_dir,
export_name + '.pb',
as_text=False)
if optimize_graph:
graph = graph_optimization(graph,
input_names=input_names,
output_names=output_names)
tf.io.write_graph(graph, export_dir, export_name + '.opt.pb')
if verbose:
print('frozen pb saved to:{}'.format(
os.path.join(export_dir, export_name)))
except Exception as e:
print(e)
finally:
K.set_session(sess_or) # rollback keras session
return
def test_SDM():
model_name = "SDM"
x, y, user_feature_columns, item_feature_columns, history_feature_list = get_xy_fd_sdm(False)
if tf.__version__ >= '2.0.0':
tf.compat.v1.disable_eager_execution()
else:
K.set_learning_phase(True)
model = SDM(user_feature_columns, item_feature_columns, history_feature_list, units=8)
# model.summary()
model.compile('adam', sampledsoftmaxloss)
check_model(model, model_name, x, y)
def test_MIND():
model_name = "MIND"
x, y, user_feature_columns, item_feature_columns = get_xy_fd(False)
K.set_learning_phase(True)
if tf.__version__ >= '2.0.0':
tf.compat.v1.disable_eager_execution()
model = MIND(user_feature_columns, item_feature_columns, num_sampled=2, user_dnn_hidden_units=(16, 4))
model.compile('adam', sampledsoftmaxloss)
check_model(model,model_name,x,y)
def main():
config = Config()
# LOAD MODEL
backend.set_learning_phase(0)
age_model = load_model(model_pth=config.age_h5_path, name="age_model")
gender_model = load_model(model_pth=config.gender_h5_path,
name="gender_model")
emotion_model = load_model(model_pth=config.expr_h5_path,
name="emotion_model")
# COMBINE
_, height, width, depth = age_model.input.shape
cb_input = layers.Input(shape=(height, width, depth))
age_outs = age_model(cb_input)
gender_outs = gender_model(cb_input)
emo_outs = emotion_model(cb_input)
merged = layers.Concatenate()([age_outs, gender_outs, emo_outs])
cb_model = models.Model(inputs=cb_input, outputs=merged)
cb_model.summary()
# SAVE MODEL
cb_model.save(config.combine_model_h5_path)
print(
colored("[INFO] Combine model is DONE, saved at </ {} /> ".format(
config.combine_model_h5_path),
color='red',
attrs=['bold']))
# FREEZE
sess = backend.get_session()
converted_output_node_names = [node.op.name for node in cb_model.outputs]
print("[INFO] in: ", cb_model.inputs) # input_1_4:0
print("[INFO] out: ", cb_model.outputs) # concatenate/concat:0
constant_graph = graph_util.convert_variables_to_constants(
sess, sess.graph.as_graph_def(), converted_output_node_names)
freeze_folder = config.freeze_folder
name = os.path.split(config.combine_model_pb_path)[-1].split('.')[0]
tf.train.write_graph(constant_graph,
freeze_folder,
'%s.pbtxt' % name,
as_text=True)
tf.train.write_graph(constant_graph,
freeze_folder,
'%s.pb' % name,
as_text=False)
print(
colored("[INFO] convert model is success, saved at </ %s/%s.pb />" %
(freeze_folder, name),
color="cyan",
attrs=['bold']))
def build_conv1fc1(self, rgb, num_classes, temp_softmax, seed, train_mode):
K.set_learning_phase(True)
# conv1_1
print("build_conv1fc1")
with tf.name_scope('mentee_conv1_1') as scope:
kernel = tf.Variable(tf.truncated_normal(
[3, 3, self.num_channels, 64],
dtype=tf.float32,
stddev=1e-2,
seed=seed),
trainable=self.trainable,
name='mentee_weights')
conv = tf.nn.conv2d(rgb, kernel, [1, 1, 1, 1], padding='SAME')
biases = tf.Variable(tf.constant(0.0, shape=[64],
dtype=tf.float32),
trainable=self.trainable,
name='mentee_biases')
out = tf.nn.bias_add(conv, biases)
# out = self.extra_regularization(out)
self.conv1_1 = tf.nn.relu(out, name=scope)
# self.conv1_1 = BatchNormalization(axis = -1, name= 'mentee_bn_conv1_1')(self.conv1_1)
self.parameters += [kernel, biases]
self.pool1 = tf.nn.max_pool(self.conv1_1,
ksize=[1, 2, 2, 1],
strides=[1, 2, 2, 1],
padding='SAME',
name='pool1')
with tf.name_scope('mentee_fc3') as scope:
shape = int(np.prod(self.pool1.get_shape()[1:]))
fc3w = tf.Variable(tf.truncated_normal([shape, num_classes],
dtype=tf.float32,
stddev=1e-2,
seed=seed),
trainable=self.trainable,
name='mentee_weights')
fc3b = tf.Variable(tf.constant(0.0,
shape=[num_classes],
dtype=tf.float32),
trainable=self.trainable,
name='mentee_biases')
pool1_flat = tf.reshape(self.pool1, [-1, shape])
self.fc3l = tf.nn.bias_add(tf.matmul(pool1_flat, fc3w), fc3b)
# self.fc3 = tf.nn.relu(fc3l)
self.parameters += [fc3w, fc3b]
self.softmax = tf.nn.softmax(self.fc3l / temp_softmax)
return self
def main():
"""Convert any Keras model to the frugally-deep model format."""
usage = 'usage: [Keras model in HDF5 format] [image output directory]'
if len(sys.argv) != 3:
print(usage)
sys.exit(1)
else:
assert K.backend() == "tensorflow"
assert K.floatx() == "float32"
assert K.image_data_format() == 'channels_last'
in_path = sys.argv[1]
out_dir = sys.argv[2]
print('loading {}'.format(in_path))
K.set_learning_phase(1)
model = load_model(in_path)
model = convert_sequential_to_model(model)
process_layers(model, out_dir)
def convert_frozen_pb_to_savedmodel(model, frozen_model_path,
frozen_model_name, savedmodel_path):
from tensorflow.python.saved_model import signature_constants
from tensorflow.python.saved_model import tag_constants
from tensorflow.python.saved_model.signature_constants import PREDICT_INPUTS
from tensorflow.python.saved_model.signature_constants import PREDICT_OUTPUTS
'''graph_pb = frozen_model_path + frozen_model_name
builder = tf.saved_model.builder.SavedModelBuilder(savedmodel_path)
with tf.gfile.GFile(graph_pb, "rb") as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())'''
builder = tf.saved_model.builder.SavedModelBuilder(savedmodel_path)
signature = {}
#with tf.Session(graph=tf.Graph()) as sess:
with K.get_session() as sess:
# name="" is important to ensure we don't get spurious prefixing
#tf.import_graph_def(graph_def, name="")
K.set_learning_phase(0)
encoder = model.get_layer('inception_v3')
encoder_out = model.get_layer('encoder_output')
graph_def = freeze_session(sess, output_names=[model.output.op.name])
tf.import_graph_def(graph_def, name="")
g = tf.get_default_graph()
#input1 = g.get_tensor_by_name(encoder.get_input_at(0).name)
input1 = g.get_tensor_by_name(model.input[0].name)
print('[INFO] Input tensor name:', input1.name)
print('[INFO] Input tensor shape:', input1.shape)
output1 = g.get_tensor_by_name(encoder_out.get_output_at(0).name)
print('[INFO] Output tensor name:', output1.name)
print('[INFO] Output tensor shape:', output1.shape)
signature[signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY] = \
tf.saved_model.signature_def_utils.predict_signature_def(
{PREDICT_INPUTS: input1},
{PREDICT_OUTPUTS: output1})
builder.add_meta_graph_and_variables(sess, [tag_constants.SERVING],
signature_def_map=signature)
builder.save()
def export_model(keras_model, export_path, model_version=0, weights_path=None):
# Start the tensorflow session
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.8,
allow_growth=False)
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))
K.set_session(sess)
K._LEARNING_PHASE = tf.constant(0)
K.set_learning_phase(0)
# Create export path if it doesn't exist
export_path = os.path.join(export_path, str(model_version))
builder = SavedModelBuilder(export_path)
# legacy_init_op = tf.group(tf.tables_initializer(), name='legacy_init_op')
# Initialize global variables and the model
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
sess.run(init_op)
# Load the model and the weights
if weights_path is not None:
keras_model.load_weights(weights_path)
if isinstance(keras_model.output, list):
output = keras_model.output[-1]
else:
output = keras_model.output
# Define prediction signature
prediction_signature = tf.saved_model.signature_def_utils.predict_signature_def(
{'image': keras_model.input}, {'prediction': output})
# Add the meta_graph and the variables to the builder
builder.add_meta_graph_and_variables(
sess, [tag_constants.SERVING],
signature_def_map={
signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
prediction_signature
})
# Save the graph
builder.save()
def save_serving_model(serving, serving_dir):
"""
keras.model을 savedModel 포맷으로 정하는 메소드
:param serving:
:param serving_dir:
:return:
"""
version_dir = get_latest_version_dir(serving_dir)
K.set_learning_phase(0)
session = K.get_session()
tf.saved_model.simple_save(session,
version_dir,
inputs={'image': serving.input},
outputs={
'visualize': serving.output[0],
'summarize': serving.output[1]
})
def _clone_and_build_model(mode,
keras_model,
custom_objects,
features=None,
labels=None):
"""Clone and build the given keras_model.
Args:
mode: training mode.
keras_model: an instance of compiled keras model.
custom_objects: Dictionary for custom objects.
features: Dict of tensors.
labels: Dict of tensors, or single tensor instance.
Returns:
The newly built model.
"""
# Set to True during training, False for inference or testing.
K.set_learning_phase(mode == model_fn_lib.ModeKeys.TRAIN)
input_tensors, target_tensors = _convert_estimator_io_to_keras(
keras_model, features, labels)
compile_clone = (mode != model_fn_lib.ModeKeys.PREDICT)
global_step = None
if compile_clone:
# Set iterations to the global step created by tf.train.create_global_step()
# which is automatically run in the estimator framework.
global_step = training_util.get_or_create_global_step()
K.track_variable(global_step)
clone = models.clone_and_build_model(
keras_model, input_tensors, target_tensors, custom_objects,
compile_clone=compile_clone,
in_place_reset=(not keras_model._is_graph_network),
optimizer_iterations=global_step)
return clone
def _clone_and_build_model(mode,
keras_model,
custom_objects,
features=None,
labels=None):
"""Clone and build the given keras_model.
Args:
mode: training mode.
keras_model: an instance of compiled keras model.
custom_objects: Dictionary for custom objects.
features: Dict of tensors.
labels: Dict of tensors, or single tensor instance.
Returns:
The newly built model.
"""
# Set to True during training, False for inference.
K.set_learning_phase(mode == model_fn_lib.ModeKeys.TRAIN)
# Get list of inputs.
if features is None:
input_tensors = None
else:
input_tensors = _create_ordered_io(keras_model,
estimator_io=features,
is_input=True)
# Get list of outputs.
if labels is None:
target_tensors = None
elif isinstance(labels, dict):
target_tensors = _create_ordered_io(keras_model,
estimator_io=labels,
is_input=False)
else:
target_tensors = [
_convert_tensor(labels)
]
if keras_model._is_graph_network:
if custom_objects:
with CustomObjectScope(custom_objects):
model = models.clone_model(keras_model, input_tensors=input_tensors)
else:
model = models.clone_model(keras_model, input_tensors=input_tensors)
else:
model = keras_model
_in_place_subclassed_model_reset(model)
if input_tensors is not None:
model._set_inputs(input_tensors)
# Compile/Build model
if mode is model_fn_lib.ModeKeys.PREDICT:
if isinstance(model, models.Sequential):
model.build()
else:
if isinstance(keras_model.optimizer, optimizers.TFOptimizer):
optimizer = keras_model.optimizer
else:
optimizer_config = keras_model.optimizer.get_config()
optimizer = keras_model.optimizer.__class__.from_config(optimizer_config)
optimizer.iterations = training_util.get_or_create_global_step()
model.compile(
optimizer,
keras_model.loss,
metrics=keras_model.metrics,
loss_weights=keras_model.loss_weights,
sample_weight_mode=keras_model.sample_weight_mode,
weighted_metrics=keras_model.weighted_metrics,
target_tensors=target_tensors)
return model
def experimental_tpu_test_loop(model,
dataset,
verbose=0,
steps=None):
"""Test loop for evaluating with TPU DistributionStrategy.
Arguments:
model: Keras Model instance.
dataset: Dataset for input data.
verbose: Integer, Verbosity mode 0 or 1.
steps: Total number of steps (batches of samples)
before declaring predictions finished.
Ignored with the default value of `None`.
Returns:
Scalar loss (if the model has a single output and no metrics)
or list of scalars (if the model has multiple outputs
and/or metrics). The attribute `model.metrics_names` will give you
the display labels for the outputs.
"""
current_strategy = model._distribution_strategy
iterator = distributed_training_utils.get_iterator(dataset, current_strategy)
scope = current_strategy.scope()
scope.__enter__()
def _per_device_eval_function(model):
model._make_eval_function()
return (model._eval_function.inputs, model._eval_function.outputs,
model._eval_function.updates_op,
model._eval_function.session_kwargs)
# TODO(priyag, sourabhbajaj): This should likely not be hardcoded here.
K.set_learning_phase(0)
def step_fn(ctx, inputs):
"""Clones the model and calls make_eval_function."""
inputs, targets = inputs
if model._compile_distribution:
distributed_training_utils. clone_model_on_replicas(
model, current_strategy,
make_callback_model=False, inputs=inputs,
targets=targets, mode=distributed_training_utils.ModeKeys.TEST)
else:
distributed_training_utils._build_distributed_network(
model, current_strategy, inputs, targets,
mode=distributed_training_utils.ModeKeys.TEST)
(grouped_inputs, grouped_outputs, grouped_updates,
grouped_session_args) = current_strategy.extended.call_for_each_replica(
_per_device_eval_function, args=(model._distributed_model_test,))
(all_inputs, all_outputs, all_updates,
all_session_args) = distributed_training_utils.unwrap_values(
current_strategy, grouped_inputs, grouped_outputs, grouped_updates,
grouped_session_args)
combined_fn = K.function(
all_inputs, all_outputs,
updates=all_updates,
name='distributed_test_function',
**all_session_args)
for label, output in zip(model.metrics_names, combined_fn.outputs):
if label == 'loss':
reduce_op = distribute_lib.get_loss_reduction()
else:
# We reduce all other metrics using mean for now. This is temporary
# workaround until new metrics are in place.
reduce_op = ds_reduce_util.ReduceOp.MEAN
ctx.set_last_step_output(label, output, reduce_op)
return combined_fn.updates_op
# Add initial dummy values for loss and other metric tensors.
initial_loop_values = {}
initial_loop_values['loss'] = constant_op.constant(1e7)
for name in model.metrics_names[1:]:
tensor = model._all_stateful_metrics_tensors[name]
initial_loop_values[name] = array_ops.zeros(tensor.shape, tensor.dtype)
# TODO(priyag): Use steps_per_run when we use new metrics as they will
# allow handling metric computation at each step using variables.
ctx = current_strategy.extended.experimental_run_steps_on_iterator(
step_fn, iterator, iterations=1,
initial_loop_values=initial_loop_values)
test_op = ctx.run_op
output_tensors = ctx.last_step_outputs
if verbose == 1:
progbar = Progbar(target=steps)
if model._compile_distribution:
distributed_training_utils._copy_weights_to_distributed_model(
model, model._distributed_model_test)
distributed_training_utils._reset_metrics(
model, model._distributed_model_test)
assert steps is not None
outs = [0.] * len(model.metrics_names)
for step in range(steps):
_, batch_outs = K.get_session().run([test_op, output_tensors])
for i, label in enumerate(model.metrics_names):
if i == 0:
# Loss is stateless metrics.
outs[i] += batch_outs[label]
else:
# For all stateful metrics, the aggregation is handled by mirrored vars.
outs[i] = batch_outs[label]
if verbose >= 1:
progbar.update(step + 1)
scope.__exit__(None, None, None)
if len(outs) >= 0:
outs[0] /= (steps)
if len(outs) == 1:
return outs[0]
return outs
def _experimental_predict_loop(model, iterator, verbose=0, steps=None):
"""Predict loop for predicting with TPU DistributionStrategy.
Arguments:
model: Keras Model instance.
iterator: Iterator for input data.
verbose: Integer, Verbosity mode 0 or 1.
steps: Total number of steps (batches of samples)
before declaring `_predict_loop` finished.
Ignored with the default value of `None`.
Returns:
Array of predictions (if the model has a single output)
or list of arrays of predictions
(if the model has multiple outputs).
"""
current_strategy = model._distribution_strategy
K.get_session().run(current_strategy.initialize())
# TODO(priyag, sourabhbajaj): This should likely not be hardcoded here.
K.set_learning_phase(0)
def _per_device_predict_function(model):
model._make_predict_function()
return (model.predict_function.inputs,
model.predict_function.outputs,
model.predict_function.updates_op,
model.predict_function.session_kwargs)
def step_fn(ctx, *inputs):
"""Clones the model and calls make_predict_function."""
# TODO(priyag, sourabhbajaj): The model gets cloned every time
# fit/test/predict is called. We should look into caching this keyed on
# input shapes.
clone_model_on_replicas(
model,
current_strategy,
make_callback_model=False,
inputs=inputs,
mode=_Mode.PREDICT)
(grouped_inputs, grouped_outputs, grouped_updates,
grouped_session_args) = current_strategy.call_for_each_replica(
_per_device_predict_function, args=(model._grouped_model_predict,))
(all_inputs, all_outputs, all_updates,
all_session_args) = distributed_training_utils.unwrap_values(
current_strategy, grouped_inputs, grouped_outputs, grouped_updates,
grouped_session_args)
combined_fn = K.function(
all_inputs, all_outputs,
updates=all_updates,
name='distributed_predict_function',
**all_session_args)
for label, output in zip(model.output_names, combined_fn.outputs):
ctx.set_last_step_output(label, output)
return combined_fn.updates_op
# Add initial dummy values for outputs.
initial_loop_values = {}
batch_dimension = distributed_training_utils.get_batch_dimension(iterator)
for name, tensor in zip(model.output_names, model.outputs):
# TODO(priyag): This is a workaround as we do not know the batch dimension
# of the model's output at this point.
shape = tensor_shape.TensorShape(tensor.shape.dims)
shape.dims = [batch_dimension] + shape.dims[1:]
initial_loop_values[name] = array_ops.zeros(shape, tensor.dtype)
with current_strategy.scope():
# TODO(priyag, sourabhbajaj): Support steps_per_run if/when we add outfeed.
ctx = current_strategy.run_steps_on_dataset(
step_fn, iterator, iterations=1,
initial_loop_values=initial_loop_values)
predict_op = ctx.run_op
output_tensors = ctx.last_step_outputs
if verbose == 1:
progbar = Progbar(target=steps)
# Copy the weights from the original model to each of the replicated models.
orig_model_weights = model.get_weights()
with current_strategy.scope():
distributed_model = current_strategy.unwrap(model._grouped_model_predict)[0]
distributed_training_utils.set_weights(
current_strategy, distributed_model, orig_model_weights)
assert steps is not None
# Since we do not know how many samples we will see, we cannot pre-allocate
# the returned Numpy arrays. Instead, we store one array per batch seen
# and concatenate them upon returning.
unconcatenated_outs = [[] for _ in model.outputs]
for step in range(steps):
_, batch_outs = K.get_session().run([predict_op, output_tensors])
# TODO(priyag): maybe need to unwrap the outputs first for MirroredStrategy.
for i, label in enumerate(model.output_names):
unconcatenated_outs[i].extend(batch_outs[label])
if verbose >= 1:
progbar.update(step + 1)
K.get_session().run(current_strategy.finalize())
if len(unconcatenated_outs) == 1:
return np.concatenate(unconcatenated_outs[0], axis=0)
return [
np.concatenate(unconcatenated_outs[i], axis=0)
for i in range(len(unconcatenated_outs))
]
def _specialize_model(self, input_specs):
"""Specialize `self.model` (a Keras model) for the given input shapes."""
# Re-create our input and output layers inside our subgraph. They will be
# attached to the true computation when we clone our model in `tpu_fn`.
K.set_learning_phase(self.execution_mode == model_fn_lib.ModeKeys.TRAIN)
# functools.partial and callable objects are not supported by tpu.rewrite
def _model_fn():
"""Compute fit/eval/predict for the TPU."""
is_training = self.execution_mode == model_fn_lib.ModeKeys.TRAIN
is_test = self.execution_mode == model_fn_lib.ModeKeys.EVAL
is_predict = self.execution_mode == model_fn_lib.ModeKeys.PREDICT
# During train/eval, we infeed our features as well as labels.
if is_training or is_test:
infeed_layers = self.model._input_layers + self.model._output_layers
else:
infeed_layers = self.model._input_layers
# Generate our infeed operation to read features & labels.
infeed_tensors = tpu_ops.infeed_dequeue_tuple(
dtypes=[spec.dtype for spec in input_specs],
shapes=[spec.shape for spec in input_specs],
name='infeed-%s' % self.execution_mode)
assert len(infeed_tensors) == len(infeed_layers), (
'Infeed inputs did not match model: %s vs %s', (infeed_layers,
infeed_tensors))
tpu_targets = []
tpu_input_map = {}
# Sort infeed outputs into inputs and labels for calling our Keras model.
for tensor, layer in zip(infeed_tensors, infeed_layers):
if layer in self.model._input_layers:
tpu_input_map[layer.name] = tensor
if layer in self.model._output_layers:
tpu_targets.append(tensor)
# Clone our CPU model, running within the TPU device context.
with TPURewriteContext(tpu_input_map):
self._cloned_model = models.clone_model(self.model)
# Create a copy of the optimizer for this graph.
if isinstance(self.model.optimizer, keras_optimizers.TFOptimizer):
cloned_optimizer = keras_optimizers.TFOptimizer(
self.model.optimizer.optimizer)
else:
logging.info('Cloning %s %s', self.model.optimizer.__class__.__name__,
self._optimizer_config)
cloned_optimizer = self.model.optimizer.__class__.from_config(
self._optimizer_config)
if is_training or is_test:
self._cloned_model.compile(
optimizer=_replicated_optimizer(cloned_optimizer),
loss=self.model.loss,
loss_weights=self.model.loss_weights,
metrics=self.model.metrics,
weighted_metrics=self.model.weighted_metrics,
target_tensors=tpu_targets,
)
# Compute our outfeed depending on the execution mode
if is_training:
self._cloned_model._make_train_function()
self._outfeed_spec = [
tensor_spec.TensorSpec(tensor.shape, tensor.dtype, tensor.name)
for tensor in self._cloned_model.train_function.outputs
]
return [
self._cloned_model.train_function.updates_op,
tpu_ops.outfeed_enqueue_tuple(
self._cloned_model.train_function.outputs,
name='outfeed-enqueue-train')
]
elif is_test:
self._cloned_model._make_test_function()
self._outfeed_spec = [
tensor_spec.TensorSpec(tensor.shape, tensor.dtype, tensor.name)
for tensor in self._cloned_model.test_function.outputs
]
return [
tpu_ops.outfeed_enqueue_tuple(
self._cloned_model.test_function.outputs,
name='outfeed-enqueue-test')
]
elif is_predict:
self._cloned_model._make_predict_function()
self._outfeed_spec = [
tensor_spec.TensorSpec(tensor.shape, tensor.dtype, tensor.name)
for tensor in self._cloned_model.predict_function.outputs
]
return [
tpu_ops.outfeed_enqueue_tuple(
self._cloned_model.predict_function.outputs,
name='outfeed-enqueue-predict',
)
]
else:
assert False, 'Unexpected execution mode: %s' % self.execution_mode
# Capture outfeed metadata computed during the rewrite.
self._outfeed_spec = None
# Generate out TPU operations using `tpu.split_compile_and_replicate`.
# `compile_op` can be used to test the TPU model compiles before execution.
# `execute op` replicates `_model_fn` `num_replicas` times, with each shard
# running on a different logical core.
compile_op, execute_op = tpu.split_compile_and_replicate(
_model_fn, inputs=[[]] * self._strategy.num_towers)
# Generate CPU side operations to enqueue features/labels and dequeue
# outputs from the model call.
infeed_op = []
outfeed_op = []
shard_infeed_tensors = []
for shard_id in range(self._strategy.num_towers):
with ops.device('/device:TPU:%d' % shard_id):
infeed_tensors = []
for spec in input_specs:
infeed_tensors.append(
array_ops.placeholder(
dtype=spec.dtype,
shape=spec.shape,
name='infeed-enqueue-%s-%d' % (spec.name, shard_id)))
shard_infeed_tensors.append(infeed_tensors)
infeed_op.append(
tpu_ops.infeed_enqueue_tuple(
infeed_tensors, [spec.shape for spec in input_specs],
name='infeed-enqueue-%s-%d' % (self.execution_mode, shard_id)))
outfeed_op.extend(
tpu_ops.outfeed_dequeue_tuple(
dtypes=[spec.dtype for spec in self._outfeed_spec],
shapes=[spec.shape for spec in self._outfeed_spec],
name='outfeed-dequeue-%s-%d' % (self.execution_mode, shard_id)))
return TPUModelOp(
compile_op,
execute_op,
infeed_tensors=shard_infeed_tensors,
infeed_op=infeed_op,
outfeed_op=outfeed_op)
def _experimental_fit_loop(
model,
iterator,
epochs=100,
verbose=1,
callbacks=None,
initial_epoch=0,
steps_per_epoch=None,
val_iterator=None,
validation_steps=None):
"""Fit loop for training with TPU DistributionStrategy.
Arguments:
model: Keras Model instance.
iterator: Iterator that returns inputs and targets
epochs: Number of times to iterate over the data
verbose: Integer, Verbosity mode, 0, 1 or 2
callbacks: List of callbacks to be called during training
initial_epoch: Epoch at which to start training
(useful for resuming a previous training run)
steps_per_epoch: Total number of steps (batches of samples)
before declaring one epoch finished and starting the
next epoch. Ignored with the default value of `None`.
val_iterator: Iterator for validation data.
validation_steps: Number of steps to run validation for
(only if doing validation from data tensors).
Ignored with the default value of `None`.
Returns:
Returns `None`.
Raises:
ValueError: in case of invalid arguments.
"""
current_strategy = model._distribution_strategy
K.get_session().run(current_strategy.initialize())
def _per_device_fit_function(model):
model._make_fit_function()
return (model._fit_function.inputs, model._fit_function.outputs,
model._fit_function.updates_op, model._fit_function.session_kwargs)
# TODO(priyag, sourabhbajaj): This should likely not be hardcoded here.
K.set_learning_phase(1)
out_labels = model.metrics_names or []
def step_fn(ctx, inputs, targets):
"""Clones the model and calls make_fit_function."""
# TODO(priyag, sourabhbajaj): The model gets cloned every time
# fit/test/predict is called. We should look into caching this keyed on
# input shapes.
clone_model_on_replicas(
model,
current_strategy,
make_callback_model=True,
inputs=inputs,
targets=targets,
mode=_Mode.TRAIN)
(grouped_inputs, grouped_outputs, grouped_updates,
grouped_session_args) = current_strategy.call_for_each_replica(
_per_device_fit_function, args=(model._grouped_model_train,))
(all_inputs, all_outputs, all_updates,
all_session_args) = distributed_training_utils.unwrap_values(
current_strategy, grouped_inputs, grouped_outputs,
grouped_updates, grouped_session_args)
combined_fn = K.function(
all_inputs,
all_outputs,
updates=all_updates,
name='distributed_fit_function',
**all_session_args)
for label, output in zip(out_labels, combined_fn.outputs):
if label == 'loss':
aggregation = distribute_lib.get_loss_reduction()
else:
# We aggregate all other metrics using mean for now. This is temporary
# workaround until new metrics are in place.
aggregation = variable_scope.VariableAggregation.MEAN
ctx.set_last_step_output(label, output, aggregation)
# TODO(priyag, sourabhbajaj): Ignoring these things from the combined_fn:
# feed_dict, session kwargs, run options, run_metadata for now. These should
# be handled appropriately
return combined_fn.updates_op
# Add initial dummy values for loss and other metric tensors.
initial_loop_values = {}
initial_loop_values['loss'] = constant_op.constant(1e7)
for name, tensor in zip(model.metrics_names[1:], model.metrics_tensors):
initial_loop_values[name] = array_ops.zeros(tensor.shape, tensor.dtype)
if steps_per_epoch is None:
raise ValueError('`steps_per_epoch` should be specified when calling '
'`fit` on the model.')
steps_per_run = K.variable(
value=min(steps_per_epoch, current_strategy.steps_per_run),
dtype='int32',
name='steps_per_run')
with current_strategy.scope():
ctx = current_strategy.run_steps_on_dataset(
step_fn, iterator, iterations=steps_per_run,
initial_loop_values=initial_loop_values)
train_op = ctx.run_op
output_tensors = ctx.last_step_outputs
do_validation = bool(validation_steps)
# Copy the weights from the original model to each of the replicated models.
orig_model_weights = model.get_weights()
with current_strategy.scope():
distributed_model = current_strategy.unwrap(model._grouped_model_train)[0]
distributed_training_utils.set_weights(
current_strategy, distributed_model, orig_model_weights)
callbacks = cbks.configure_callbacks(
callbacks,
model,
do_validation=do_validation,
val_inputs=None,
val_targets=None,
epochs=epochs,
steps_per_epoch=steps_per_epoch,
verbose=verbose)
# Calculate the steps each time on the device.
steps_to_run = [current_strategy.steps_per_run] * (
steps_per_epoch // current_strategy.steps_per_run)
if steps_per_epoch % current_strategy.steps_per_run:
steps_to_run.append(steps_per_epoch % current_strategy.steps_per_run)
callbacks.on_train_begin()
for epoch in range(initial_epoch, epochs):
callbacks.on_epoch_begin(epoch)
epoch_logs = {}
step_index = 0
prev_step_count = None
for step_count in steps_to_run:
batch_logs = {'batch': step_index, 'size': 1, 'num_steps': step_count}
callbacks.on_batch_begin(step_index, batch_logs)
if prev_step_count is None or step_count != prev_step_count:
steps_per_run.load(step_count, K.get_session())
prev_step_count = step_count
try:
_, outputs = K.get_session().run([train_op, output_tensors])
except errors.OutOfRangeError:
logging.warning('Your dataset iterator ran out of data; '
'interrupting training. Make sure that your dataset '
'can generate at least `steps_per_epoch * epochs` '
'batches (in this case, %d batches).' %
steps_per_epoch * epochs)
break
batch_logs.update(outputs)
callbacks.on_batch_end(step_index, batch_logs)
step_index = step_index + step_count
if callbacks.model.stop_training:
break
if do_validation:
logging.info('Running validation at fit epoch: %s', epoch)
# Since we create a new clone from the original model we need to copy
# the weights back to the original model before we can run validation.
with current_strategy.scope():
updated_weights = current_strategy.unwrap(
model._grouped_model_train)[0].get_weights()
model.set_weights(updated_weights)
val_outs = _experimental_test_loop(
model,
val_iterator,
steps=validation_steps,
verbose=verbose,
initialize_finalize_strategy=False)
if not isinstance(val_outs, list):
val_outs = [val_outs]
# Same labels assumed.
for label, val_out in zip(out_labels, val_outs):
epoch_logs['val_' + label] = val_out
callbacks.on_epoch_end(epoch, epoch_logs)
if callbacks.model.stop_training:
break
callbacks.on_train_end()
# Copy the weights back from the replicated model to the original model.
with current_strategy.scope():
updated_weights = current_strategy.unwrap(
model._grouped_model_train)[0].get_weights()
model.set_weights(updated_weights)
K.get_session().run(current_strategy.finalize())
return model.history
def _experimental_test_loop(model, iterator, verbose=0, steps=None,
initialize_finalize_strategy=True):
"""Test loop for evaluating with TPU DistributionStrategy.
Arguments:
model: Keras Model instance.
iterator: Iterator for input data.
verbose: Integer, Verbosity mode 0 or 1.
steps: Total number of steps (batches of samples)
before declaring predictions finished.
Ignored with the default value of `None`.
initialize_finalize_strategy: Should the strategy initialize and finalize
functions be called.
Returns:
Scalar loss (if the model has a single output and no metrics)
or list of scalars (if the model has multiple outputs
and/or metrics). The attribute `model.metrics_names` will give you
the display labels for the outputs.
"""
current_strategy = model._distribution_strategy
if initialize_finalize_strategy:
K.get_session().run(current_strategy.initialize())
def _per_device_eval_function(model):
model._make_eval_function()
return (model._eval_function.inputs, model._eval_function.outputs,
model._eval_function.updates_op,
model._eval_function.session_kwargs)
# TODO(priyag, sourabhbajaj): This should likely not be hardcoded here.
K.set_learning_phase(0)
def step_fn(ctx, inputs, targets):
"""Clones the model and calls make_eval_function."""
# TODO(priyag, sourabhbajaj): The model gets cloned every time
# fit/test/predict is called. We should look into caching this keyed on
# input shapes.
clone_model_on_replicas(
model,
current_strategy,
make_callback_model=False,
inputs=inputs,
targets=targets,
mode=_Mode.TEST)
(grouped_inputs, grouped_outputs, grouped_updates,
grouped_session_args) = current_strategy.call_for_each_replica(
_per_device_eval_function, args=(model._grouped_model_test,))
(all_inputs, all_outputs, all_updates,
all_session_args) = distributed_training_utils.unwrap_values(
current_strategy, grouped_inputs, grouped_outputs, grouped_updates,
grouped_session_args)
combined_fn = K.function(
all_inputs, all_outputs,
updates=all_updates,
name='distributed_test_function',
**all_session_args)
for label, output in zip(model.metrics_names, combined_fn.outputs):
if label == 'loss':
aggregation = distribute_lib.get_loss_reduction()
else:
# We aggregate all other metrics using mean for now. This is temporary
# workaround until new metrics are in place.
aggregation = variable_scope.VariableAggregation.MEAN
ctx.set_last_step_output(label, output, aggregation)
return combined_fn.updates_op
# Add initial dummy values for loss and other metric tensors.
initial_loop_values = {}
initial_loop_values['loss'] = constant_op.constant(1e7)
for name, tensor in zip(model.metrics_names[1:], model.metrics_tensors):
initial_loop_values[name] = array_ops.zeros(tensor.shape, tensor.dtype)
with current_strategy.scope():
# TODO(priyag): Use steps_per_run when we use new metrics as they will
# allow handling metric computation at each step using variables.
ctx = current_strategy.run_steps_on_dataset(
step_fn, iterator, iterations=1,
initial_loop_values=initial_loop_values)
test_op = ctx.run_op
output_tensors = ctx.last_step_outputs
if verbose == 1:
progbar = Progbar(target=steps)
# Copy the weights from the original model to each of the replicated models.
orig_model_weights = model.get_weights()
with current_strategy.scope():
distributed_model = current_strategy.unwrap(model._grouped_model_test)[0]
distributed_training_utils.set_weights(
current_strategy, distributed_model, orig_model_weights)
assert steps is not None
outs = [0.] * len(model.metrics_names)
for step in range(steps):
_, batch_outs = K.get_session().run([test_op, output_tensors])
for i, label in enumerate(model.metrics_names):
outs[i] += batch_outs[label]
if verbose >= 1:
progbar.update(step + 1)
for i in range(len(outs)):
outs[i] /= (steps)
if initialize_finalize_strategy:
K.get_session().run(current_strategy.finalize())
if len(outs) == 1:
return outs[0]
return outs
def model_iteration(model,
data,
steps_per_epoch=None,
epochs=1,
verbose=1,
callbacks=None,
validation_data=None,
validation_steps=None,
validation_freq=1,
class_weight=None,
max_queue_size=10,
workers=1,
use_multiprocessing=False,
shuffle=False,
initial_epoch=0,
mode=ModeKeys.TRAIN,
batch_size=None,
steps_name='steps',
**kwargs):
"""Loop function for arrays of data with modes TRAIN/TEST/PREDICT.
Arguments:
model: Keras Model instance.
data: Either a tuple of NumPy/Tensor inputs (i.e. `(x,)` or `(x, y)` or
`(x, y, sample_weights)`) or a generator or
`keras.utils.data_utils.Sequence` object or Eager Iterator or Dataset.
steps_per_epoch: Total number of steps (batches of samples) before
declaring one epoch finished and starting the next epoch. Ignored with
the default value of `None`.
epochs: Number of times to iterate over the data.
verbose: Verbosity mode, 0, 1 or 2.
callbacks: List of callbacks to be called during training.
validation_data: Either a tuple of NumPy/Tensor inputs (i.e. `(x,)` or
`(x, y)` or `(x, y, sample_weights)`) or a generator or
`keras.utils.data_utils.Sequence` object or Eager Iterator or Dataset.
validation_steps: Total number of steps (batches of samples) before
declaring validation finished.
validation_freq: Only relevant if validation data is provided. Integer or
`collections.Container` instance (e.g. list, tuple, etc.). If an
integer, specifies how many training epochs to run before a new
validation run is performed, e.g. `validation_freq=2` runs
validation every 2 epochs. If a Container, specifies the epochs on
which to run validation, e.g. `validation_freq=[1, 2, 10]` runs
validation at the end of the 1st, 2nd, and 10th epochs.
class_weight: Dictionary mapping class indices to a weight for the class.
max_queue_size: Integer. Maximum size for the generator queue. If
unspecified, `max_queue_size` will default to 10.
workers: Integer. Maximum number of processes to spin up when using
process-based threading. If unspecified, `workers` will default to 1. If
0, will execute the generator on the main thread.
use_multiprocessing: Boolean. If `True`, use process-based threading. If
unspecified, `use_multiprocessing` will default to `False`. Note that
because this implementation relies on multiprocessing, you should not
pass non-picklable arguments to the generator as they can't be passed
easily to children processes.
shuffle: Boolean. Whether to shuffle the order of the batches at the
beginning of each epoch. Only used with instances of `Sequence`
(`keras.utils.Sequence`). Has no effect when `steps_per_epoch` is not
`None`.
initial_epoch: Epoch at which to start training (useful for resuming a
previous training run).
mode: One of ModeKeys.TRAIN/ModeKeys.TEST/ModeKeys.PREDICT.
batch_size: Integer batch size or None if unknown. Will only be used if
`data` is in NumPy/Tensor format.
steps_name: The string name of the steps argument, either `steps`,
`validation_steps`, or `steps_per_epoch`. Only used for error message
formatting.
**kwargs: Additional arguments for backwards compatibility. `steps` is
accepted as an alias for `steps_per_epoch`.
Returns:
- In TRAIN mode: `History` object.
- In TEST mode: Evaluation metrics.
- In PREDICT mode: Outputs of the Model called on inputs.
Raises:
ValueError: in case of invalid arguments.
"""
if 'steps' in kwargs:
steps_per_epoch = kwargs['steps']
# Determine the number of steps per epoch and whether we should reset the
# dataset at the end of each epoch.
reset_dataset_after_each_epoch = False
original_dataset = None
is_dataset = isinstance(data, (dataset_ops.DatasetV2, dataset_ops.DatasetV1))
if is_dataset:
original_dataset = data
if steps_per_epoch is None:
reset_dataset_after_each_epoch = True
steps_per_epoch = training_utils.infer_steps_for_dataset(
data, steps_per_epoch, epochs=epochs, steps_name=steps_name)
# Convert to a format that supports `next(generator)`.
generator, steps_per_epoch = convert_to_generator_like(
data,
steps_per_epoch=steps_per_epoch,
batch_size=batch_size,
epochs=epochs - initial_epoch,
shuffle=shuffle)
do_validation = validation_data is not None
should_set_learning_phase = context.executing_eagerly() and model.run_eagerly
is_sequence = isinstance(generator, data_utils.Sequence)
_validate_arguments(is_sequence, is_dataset, use_multiprocessing, workers,
steps_per_epoch, validation_data, validation_steps, mode,
kwargs)
batch_function = _make_execution_function(
model, mode, class_weight=class_weight)
# Create the queue for the generator.
output_generator, enqueuer = _make_enqueued_generator(
generator,
workers=workers,
use_multiprocessing=use_multiprocessing,
max_queue_size=max_queue_size,
shuffle=shuffle)
num_samples_or_steps, use_steps = _get_num_samples_or_steps(
data, steps_per_epoch)
count_mode = 'steps' if use_steps else 'samples'
callbacks = cbks.configure_callbacks(
callbacks,
model,
do_validation=do_validation,
epochs=epochs,
steps_per_epoch=steps_per_epoch,
batch_size=batch_size,
samples=num_samples_or_steps,
verbose=0, # Handle ProgBar as part of Callbacks once hooks are ready.
mode=mode)
# TODO(omalleyt): Handle ProgBar as part of Callbacks once hooks are ready.
progbar = training_utils.get_progbar(model, count_mode)
progbar.params = callbacks.params
progbar.params['verbose'] = verbose
if mode == ModeKeys.PREDICT:
aggregator = training_utils.OutputsAggregator(True, steps_per_epoch)
else:
aggregator = training_utils.MetricsAggregator(True, steps_per_epoch)
if should_set_learning_phase:
old_learning_phase = backend.learning_phase()
backend.set_learning_phase(1 if mode == ModeKeys.TRAIN else 0)
callbacks.model.stop_training = False
callbacks._call_begin_hook(mode)
progbar.on_train_begin()
for epoch in range(initial_epoch, epochs):
if callbacks.model.stop_training:
break
# Setup work for each epoch.
model.reset_metrics()
epoch_logs = {}
if mode == ModeKeys.TRAIN:
callbacks.on_epoch_begin(epoch, epoch_logs)
progbar.on_epoch_begin(epoch, epoch_logs)
if steps_per_epoch is None:
# Loop over dataset until `OutOfRangeError` is raised.
target_steps = np.inf
else:
# Loop over dataset for the specified number of steps.
target_steps = steps_per_epoch
step = 0
while step < target_steps:
batch_data = _get_next_batch(output_generator, mode)
if batch_data is None:
if not is_dataset:
# We ran out of batches while the user passed an iterator (legacy).
logging.warning(
'Your dataset iterator ran out of data; '
'interrupting training. Make sure that your iterator '
'can generate at least `%s * epochs` '
'batches (in this case, %d batches). You may need to'
'use the repeat() function when building your '
'dataset.' % (steps_name, steps_per_epoch * epochs))
callbacks.model.stop_training = True
else:
# The dataset passed by the user ran out of batches.
# Now we know the cardinality of the dataset.
# assert steps_per_epoch is None
if step > 0:
steps_per_epoch = step
aggregator.num_samples_or_steps = steps_per_epoch
progbar.params['steps'] = steps_per_epoch
progbar.progbar.target = steps_per_epoch
break
# `batch_size` used for validation data if validation
# data is NumPy/EagerTensors.
batch_size = int(nest.flatten(batch_data)[0].shape[0])
# Callbacks batch begin.
batch_logs = {'batch': step, 'size': batch_size}
callbacks._call_batch_hook(mode, 'begin', step, batch_logs)
progbar.on_batch_begin(step, batch_logs)
batch_outs = batch_function(*batch_data)
if not isinstance(batch_outs, list):
batch_outs = [batch_outs]
# Aggregate results.
if step == 0:
aggregator.create(batch_outs)
aggregator.aggregate(batch_outs)
# Callbacks batch end.
batch_logs = cbks.make_logs(model, batch_logs, batch_outs, mode)
callbacks._call_batch_hook(mode, 'end', step, batch_logs)
progbar.on_batch_end(step, batch_logs)
step += 1
if callbacks.model.stop_training:
break
aggregator.finalize()
results = aggregator.results
epoch_logs = cbks.make_logs(model, epoch_logs, results, mode)
if len(results) == 1:
results = results[0]
# Run the test loop every epoch during training.
if (do_validation and
training_utils.should_run_validation(validation_freq, epoch) and
not callbacks.model.stop_training):
val_results = model_iteration(
model,
validation_data,
steps_per_epoch=validation_steps,
batch_size=batch_size,
class_weight=class_weight,
workers=workers,
use_multiprocessing=use_multiprocessing,
max_queue_size=max_queue_size,
callbacks=callbacks,
verbose=0,
mode=ModeKeys.TEST,
steps_name='validation_steps')
if not isinstance(val_results, list):
val_results = [val_results]
epoch_logs = cbks.make_logs(
model, epoch_logs, val_results, mode, prefix='val_')
if mode == ModeKeys.TRAIN:
# Epochs only apply to `fit`.
callbacks.on_epoch_end(epoch, epoch_logs)
progbar.on_epoch_end(epoch, epoch_logs)
# Recreate dataset iterator for the next epoch.
if reset_dataset_after_each_epoch and epoch < epochs - 1:
generator = dataset_ops.make_one_shot_iterator(original_dataset)
callbacks._call_end_hook(mode)
if enqueuer is not None:
enqueuer.stop()
if should_set_learning_phase:
backend.set_learning_phase(old_learning_phase)
if mode == ModeKeys.TRAIN:
return model.history
return results
def model_iteration(model,
data,
steps_per_epoch=None,
epochs=1,
verbose=1,
callbacks=None,
validation_data=None,
validation_steps=None,
class_weight=None,
max_queue_size=10,
workers=1,
use_multiprocessing=False,
shuffle=True,
initial_epoch=0,
mode='train',
batch_size=None,
**kwargs):
"""Loop function for arrays of data with modes 'train'/'test'/'predict'.
Arguments:
model: Keras Model instance.
data: Either a tuple of NumPy/Tensor inputs (i.e. `(x,)` or `(x, y)` or
`(x, y, sample_weights)`) or a generator or
`keras.utils.data_utils.Sequence` object or Eager Iterator or Dataset.
steps_per_epoch: Total number of steps (batches of samples) before
declaring one epoch finished and starting the next epoch. Ignored with
the default value of `None`.
epochs: Number of times to iterate over the data.
verbose: Verbosity mode, 0, 1 or 2.
callbacks: List of callbacks to be called during training.
validation_data: Either a tuple of NumPy/Tensor inputs (i.e. `(x,)` or
`(x, y)` or `(x, y, sample_weights)`) or a generator or
`keras.utils.data_utils.Sequence` object or Eager Iterator or Dataset.
validation_steps: Total number of steps (batches of samples) before
declaring validation finished.
class_weight: Dictionary mapping class indices to a weight for the class.
max_queue_size: Integer. Maximum size for the generator queue. If
unspecified, `max_queue_size` will default to 10.
workers: Integer. Maximum number of processes to spin up when using
process-based threading. If unspecified, `workers` will default to 1. If
0, will execute the generator on the main thread.
use_multiprocessing: Boolean. If `True`, use process-based threading. If
unspecified, `use_multiprocessing` will default to `False`. Note that
because this implementation relies on multiprocessing, you should not
pass non-picklable arguments to the generator as they can't be passed
easily to children processes.
shuffle: Boolean. Whether to shuffle the order of the batches at the
beginning of each epoch. Only used with instances of `Sequence`
(`keras.utils.Sequence`). Has no effect when `steps_per_epoch` is not
`None`.
initial_epoch: Epoch at which to start training (useful for resuming a
previous training run).
mode: One of 'train'/'test'/'predict'.
batch_size: Integer batch size or None if unknown. Will only be used if
`data` is in NumPy/Tensor format.
**kwargs: Additional arguments for backwards compatibility. `steps` is
accepted as an alias for `steps_per_epoch`.
Returns:
- In 'train' mode: `History` object.
- In 'test' mode: Evaluation metrics.
- In 'predict' mode: Outputs of the Model called on inputs.
Raises:
ValueError: in case of invalid arguments.
"""
if 'steps' in kwargs:
steps_per_epoch = kwargs['steps']
# Convert to a format that supports `next(generator)`.
generator, steps_per_epoch = convert_to_generator_like(
data,
steps_per_epoch=steps_per_epoch,
batch_size=batch_size,
epochs=epochs - initial_epoch,
shuffle=shuffle)
do_validation = validation_data is not None
should_set_learning_phase = context.executing_eagerly() and model.run_eagerly
is_sequence = isinstance(generator, data_utils.Sequence)
_validate_arguments(is_sequence, use_multiprocessing, workers,
steps_per_epoch, validation_data, validation_steps, mode,
kwargs)
batch_function = _make_execution_function(
model, mode, class_weight=class_weight)
# Create the queue for the generator.
output_generator, enqueuer = _make_enqueued_generator(
generator,
workers=workers,
use_multiprocessing=use_multiprocessing,
max_queue_size=max_queue_size,
shuffle=shuffle)
num_samples_or_steps, use_steps = _get_num_samples_or_steps(
data, steps_per_epoch)
count_mode = 'steps' if use_steps else 'samples'
callbacks = cbks.configure_callbacks(
callbacks,
model,
do_validation=do_validation,
epochs=epochs,
steps_per_epoch=steps_per_epoch,
batch_size=batch_size,
samples=num_samples_or_steps,
verbose=0, # Handle ProgBar as part of Callbacks once hooks are ready.
mode=mode)
# TODO(omalleyt): Handle ProgBar as part of Callbacks once hooks are ready.
progbar = training_utils.get_progbar(model, count_mode)
progbar.params = callbacks.params
progbar.params['verbose'] = verbose
if mode == 'predict':
aggregator = training_utils.OutputsAggregator(True, steps_per_epoch)
else:
aggregator = training_utils.MetricsAggregator(True, steps_per_epoch)
if should_set_learning_phase:
old_learning_phase = backend.learning_phase()
backend.set_learning_phase(1 if mode == 'train' else 0)
callbacks.model.stop_training = False
callbacks._call_begin_hook(mode)
progbar.on_train_begin()
for epoch in range(initial_epoch, epochs):
if callbacks.model.stop_training:
break
# Setup work for each epoch.
model.reset_metrics()
epoch_logs = {}
callbacks.on_epoch_begin(epoch, epoch_logs, mode=mode)
progbar.on_epoch_begin(epoch, epoch_logs)
for step in range(steps_per_epoch):
batch_data = _get_next_batch(output_generator, mode)
if batch_data is None:
callbacks.model.stop_training = True
break
# `batch_size` used for validation data if validation
# data is NumPy/EagerTensors.
batch_size = int(nest.flatten(batch_data)[0].shape[0])
# Callbacks batch begin.
batch_logs = {'batch': step, 'size': batch_size}
callbacks._call_batch_hook(mode, 'begin', step, batch_logs)
progbar.on_batch_begin(step, batch_logs)
batch_outs = batch_function(*batch_data)
if not isinstance(batch_outs, list):
batch_outs = [batch_outs]
# Aggregate results.
if step == 0:
aggregator.create(batch_outs)
aggregator.aggregate(batch_outs)
# Callbacks batch end.
batch_logs.update(training_utils.make_logs(model, batch_outs, mode))
callbacks._call_batch_hook(mode, 'end', step, batch_logs)
progbar.on_batch_end(step, batch_logs)
if callbacks.model.stop_training:
break
aggregator.finalize()
results = aggregator.results
epoch_logs.update(training_utils.make_logs(model, results, mode))
if len(results) == 1:
results = results[0]
# Run the test loop every epoch during training.
if do_validation and not callbacks.model.stop_training:
val_results = model_iteration(
model,
validation_data,
steps_per_epoch=validation_steps,
batch_size=batch_size,
class_weight=class_weight,
workers=workers,
use_multiprocessing=use_multiprocessing,
max_queue_size=max_queue_size,
mode='test')
if not isinstance(val_results, list):
val_results = [val_results]
epoch_logs.update(
training_utils.make_logs(model, val_results, mode, prefix='val_'))
callbacks.on_epoch_end(epoch, epoch_logs, mode=mode)
progbar.on_epoch_end(epoch, epoch_logs)
callbacks._call_end_hook(mode)
if enqueuer is not None:
enqueuer.stop()
if should_set_learning_phase:
backend.set_learning_phase(old_learning_phase)
if mode == 'train':
return model.history
return results
def _experimental_fit_loop(
model,
iterator,
epochs=100,
initial_epoch=0,
steps_per_epoch=None):
"""fit function when using TPU DistributionStrategy for training.
Arguments:
model: Keras Model instance.
iterator: Iterator that returns inputs and targets
epochs: Number of times to iterate over the data
initial_epoch: Epoch at which to start training
(useful for resuming a previous training run)
steps_per_epoch: Total number of steps (batches of samples)
before declaring one epoch finished and starting the
next epoch. Ignored with the default value of `None`.
Returns:
Returns `None`.
Raises:
ValueError: in case of invalid arguments.
"""
current_strategy = model._distribution_strategy
# TODO(priyag): Add validation that shapes are fully defined for TPU case.
# TODO(priyag, sourabhbajaj): This should be moved into a callback instead.
K.get_session().run(current_strategy.initialize())
def _per_device_train_function(model):
model._make_train_function()
return (model.train_function.inputs,
model.train_function.outputs,
model.train_function.updates_op,
model.train_function.session_kwargs)
# TODO(priyag, sourabhbajaj): This should likely not be hardcoded here.
K.set_learning_phase(1)
def step_fn(ctx, inputs, targets):
"""Clones the model and calls make_train_function."""
# TODO(priyag, sourabhbajaj): Should cache this keyed on input shapes.
clone_model_on_towers(
model,
current_strategy,
make_callback_model=True,
inputs=inputs,
targets=targets)
(grouped_inputs, grouped_outputs, grouped_updates,
grouped_session_args) = current_strategy.call_for_each_tower(
_per_device_train_function, model._grouped_model)
(all_inputs, all_outputs, all_updates,
all_session_args) = distributed_training_utils.unwrap_values(
current_strategy, grouped_inputs, grouped_outputs,
grouped_updates, grouped_session_args, with_loss_tensor=True)
combined_fn = K.Function(
all_inputs, all_outputs,
updates=all_updates,
name='distributed_train_function',
**all_session_args)
# TODO(priyag, sourabhbajaj): Perhaps the aggregation type needs to be
# something else for different outputs.
out_labels = model.metrics_names or []
for label, output in zip(out_labels, combined_fn.outputs):
ctx.set_last_step_output(label, output,
aggregation=distribute_lib.get_loss_reduction())
# TODO(priyag, sourabhbajaj): Ignoring these things from the combined_fn:
# feed_dict, session kwargs, run options, run_metadata for now. These should
# be handled appropriately
return combined_fn.updates_op
# Add initial dummy values for loss and other metric tensors.
initial_loop_values = {}
initial_loop_values['loss'] = constant_op.constant(1e7)
for name, tensor in zip(model.metrics_names[1:], model.metrics_tensors):
initial_loop_values[name] = array_ops.zeros(tensor.shape, tensor.dtype)
with current_strategy.scope():
# TODO(priyag, sourabhbajaj): Adjust steps_per_run appropriately based on
# steps_per_epoch and number of epochs.
ctx = current_strategy.run_steps_on_dataset(
step_fn, iterator, iterations=current_strategy.steps_per_run,
initial_loop_values=initial_loop_values)
train_op = ctx.run_op
output_tensors = ctx.last_step_outputs
# Copy the weights from the original model to each of the replicated models.
orig_model_weights = model.get_weights()
with current_strategy.scope():
distributed_model = current_strategy.unwrap(model._grouped_model)[0]
distributed_training_utils.set_weights(
current_strategy, distributed_model, orig_model_weights)
assert steps_per_epoch is not None
# TODO(priyag, sourabhbajaj): Add callbacks support.
# TODO(priyag, sourabhbajaj): Add validation.
for epoch in range(initial_epoch, epochs):
for step_index in range(
0, steps_per_epoch, current_strategy.steps_per_run):
try:
_, outs = K.get_session().run([train_op, output_tensors])
# TODO(priyag, sourabhbajaj): Remove this logging in favor of proper
# summaries through callbacks.
print('Epoch: {}, step_index: {}, loss: {}'.format(
epoch, step_index, outs['loss']))
for label, out in outs.items():
print(label, ': ', out)
except errors.OutOfRangeError:
logging.warning('Your dataset iterator ran out of data; '
'interrupting training. Make sure that your dataset '
'can generate at least `steps_per_epoch * epochs` '
'batches (in this case, %d batches).' %
steps_per_epoch * epochs)
break
# Copy the weights back from the replicated model to the original model.
with current_strategy.scope():
updated_weights = current_strategy.unwrap(
model._grouped_model)[0].get_weights()
model.set_weights(updated_weights)
K.get_session().run(current_strategy.finalize())
def _export_mode(
mode, has_saved_vars, builder, model, custom_objects, checkpoint_path):
"""Export a model, and optionally save new vars from the clone model.
Args:
mode: A `tf.estimator.ModeKeys` string.
has_saved_vars: A `boolean` indicating whether the SavedModel has already
exported variables.
builder: A `SavedModelBuilder` object.
model: A `tf.keras.Model` object.
custom_objects: A dictionary mapping string names to custom classes
or functions.
checkpoint_path: String path to checkpoint.
Raises:
ValueError: If the train/eval mode is being exported, but the model does
not have an optimizer.
"""
compile_clone = (mode != model_fn_lib.ModeKeys.PREDICT)
if compile_clone and not model.optimizer:
raise ValueError(
'Model does not have an optimizer. Cannot export mode %s' % mode)
model_graph = ops.get_default_graph()
with ops.Graph().as_default() as g:
K.set_learning_phase(mode == model_fn_lib.ModeKeys.TRAIN)
# Clone the model into blank graph. This will create placeholders for inputs
# and targets.
clone = models_lib.clone_and_build_model(
model, custom_objects=custom_objects, compile_clone=compile_clone)
# Make sure that iterations variable is added to the global step collection,
# to ensure that, when the SavedModel graph is loaded, the iterations
# variable is returned by `tf.train.get_global_step()`. This is required for
# compatibility with the SavedModelEstimator.
if compile_clone:
g.add_to_collection(ops.GraphKeys.GLOBAL_STEP, clone.optimizer.iterations)
# Extract update and train ops from train/test/predict functions.
train_op = None
if mode == model_fn_lib.ModeKeys.TRAIN:
clone._make_train_function()
train_op = clone.train_function.updates_op
elif mode == model_fn_lib.ModeKeys.EVAL:
clone._make_test_function()
else:
clone._make_predict_function()
g.get_collection_ref(ops.GraphKeys.UPDATE_OPS).extend(clone.state_updates)
clone_var_list = checkpointable_utils.named_saveables(clone)
with session.Session().as_default():
if has_saved_vars:
# Confirm all variables in the clone have an entry in the checkpoint.
status = clone.load_weights(checkpoint_path)
status.assert_existing_objects_matched()
else:
# Confirm that variables between the clone and model match up exactly,
# not counting optimizer objects. Optimizer objects are ignored because
# if the model has not trained, the slot variables will not have been
# created yet.
# TODO(b/113179535): Replace with checkpointable equivalence.
_assert_same_non_optimizer_objects(model, model_graph, clone, g)
# TODO(b/113178242): Use value transfer for checkpointable objects.
clone.load_weights(checkpoint_path)
# Add graph and variables to SavedModel.
# TODO(b/113134168): Switch to add_meta_graph_and_variables.
clone.save_weights(checkpoint_path, save_format='tf', overwrite=True)
builder._has_saved_variables = True
# Add graph to the SavedModel builder.
builder.add_meta_graph(
model_fn_lib.EXPORT_TAG_MAP[mode],
signature_def_map=_create_signature_def_map(clone, mode),
saver=saver_lib.Saver(clone_var_list),
init_op=variables.local_variables_initializer(),
train_op=train_op)
return None
def experimental_tpu_fit_loop(model,
dataset,
epochs=100,
verbose=1,
callbacks=None,
initial_epoch=0,
steps_per_epoch=None,
val_dataset=None,
validation_steps=None,
validation_freq=1):
"""Fit loop for training with TPU DistributionStrategy.
Arguments:
model: Keras Model instance.
dataset: Dataset that returns inputs and targets
epochs: Number of times to iterate over the data
verbose: Integer, Verbosity mode, 0, 1 or 2
callbacks: List of callbacks to be called during training
initial_epoch: Epoch at which to start training
(useful for resuming a previous training run)
steps_per_epoch: Total number of steps (batches of samples)
before declaring one epoch finished and starting the
next epoch. Ignored with the default value of `None`.
val_dataset: Dataset for validation data.
validation_steps: Number of steps to run validation for
(only if doing validation from data tensors).
Ignored with the default value of `None`.
validation_freq: Only relevant if validation data is provided. Integer or
`collections.Container` instance (e.g. list, tuple, etc.). If an
integer, specifies how many training epochs to run before a new
validation run is performed, e.g. `validation_freq=2` runs
validation every 2 epochs. If a Container, specifies the epochs on
which to run validation, e.g. `validation_freq=[1, 2, 10]` runs
validation at the end of the 1st, 2nd, and 10th epochs.
Returns:
Returns `None`.
Raises:
ValueError: in case of invalid arguments.
"""
# TODO(fchollet): add support for `steps_per_epoch=None` in TPU loops.
current_strategy = model._distribution_strategy
iterator = distributed_training_utils.get_iterator(dataset, current_strategy)
scope = current_strategy.scope()
scope.__enter__()
def _per_device_fit_function(model):
model._make_fit_function()
return (model._fit_function.inputs, model._fit_function.outputs,
model._fit_function.updates_op, model._fit_function.session_kwargs)
# TODO(priyag, sourabhbajaj): This should likely not be hardcoded here.
K.set_learning_phase(1)
out_labels = model.metrics_names or []
def step_fn(ctx, inputs):
"""Clones the model and calls make_fit_function."""
inputs, targets = inputs
if model._compile_distribution:
distributed_training_utils.clone_model_on_replicas(
model, current_strategy,
make_callback_model=True, inputs=inputs,
targets=targets, mode=distributed_training_utils.ModeKeys.TRAIN)
else:
distributed_training_utils._build_distributed_network(
model, current_strategy, inputs,
targets, mode=distributed_training_utils.ModeKeys.TRAIN)
(grouped_inputs, grouped_outputs, grouped_updates,
grouped_session_args) = current_strategy.extended.call_for_each_replica(
_per_device_fit_function, args=(model._distributed_model_train,))
(all_inputs, all_outputs, all_updates,
all_session_args) = distributed_training_utils.unwrap_values(
current_strategy, grouped_inputs, grouped_outputs,
grouped_updates, grouped_session_args)
combined_fn = K.function(
all_inputs,
all_outputs,
updates=all_updates,
name='distributed_fit_function',
**all_session_args)
for label, output in zip(out_labels, combined_fn.outputs):
if label == 'loss':
reduce_op = distribute_lib.get_loss_reduction()
else:
# We reduce all other metrics using mean for now. This is temporary
# workaround until new metrics are in place.
reduce_op = ds_reduce_util.ReduceOp.MEAN
ctx.set_last_step_output(label, output, reduce_op)
# TODO(priyag, sourabhbajaj): Ignoring these things from the combined_fn:
# feed_dict, session kwargs, run options, run_metadata for now. These should
# be handled appropriately
return combined_fn.updates_op
# Add initial dummy values for loss and other metric tensors.
initial_loop_values = {}
initial_loop_values['loss'] = constant_op.constant(1e7)
for name in model.metrics_names[1:]:
tensor = model._all_stateful_metrics_tensors[name]
initial_loop_values[name] = array_ops.zeros(tensor.shape, tensor.dtype)
if steps_per_epoch is None:
raise ValueError('`steps_per_epoch` should be specified when calling '
'`fit` on the model.')
steps_per_run = K.variable(
value=min(steps_per_epoch, current_strategy.extended.steps_per_run),
dtype='int32',
name='steps_per_run')
ctx = current_strategy.extended.experimental_run_steps_on_iterator(
step_fn, iterator, iterations=steps_per_run,
initial_loop_values=initial_loop_values)
train_op = ctx.run_op
output_tensors = ctx.last_step_outputs
do_validation = bool(validation_steps)
if model._compile_distribution:
distributed_training_utils._copy_weights_to_distributed_model(
model, model._distributed_model_train)
callbacks = cbks.configure_callbacks(
callbacks,
model,
do_validation=do_validation,
epochs=epochs,
steps_per_epoch=steps_per_epoch,
verbose=verbose)
# Calculate the steps each time on the device.
steps_to_run = [current_strategy.extended.steps_per_run] * (
steps_per_epoch // current_strategy.extended.steps_per_run)
if steps_per_epoch % current_strategy.extended.steps_per_run:
steps_to_run.append(
steps_per_epoch % current_strategy.extended.steps_per_run)
callbacks.on_train_begin()
for epoch in range(initial_epoch, epochs):
distributed_training_utils._reset_metrics(
model, model._distributed_model_train)
callbacks.on_epoch_begin(epoch)
epoch_logs = {}
step_index = 0
prev_step_count = None
for step_count in steps_to_run:
batch_logs = {'batch': step_index, 'size': 1, 'num_steps': step_count}
callbacks.on_batch_begin(step_index, batch_logs)
if prev_step_count is None or step_count != prev_step_count:
steps_per_run.load(step_count, K.get_session())
prev_step_count = step_count
try:
_, outputs = K.get_session().run([train_op, output_tensors])
except errors.OutOfRangeError:
logging.warning('Your dataset iterator ran out of data; '
'interrupting training. Make sure that your dataset '
'can generate at least `steps_per_epoch * epochs` '
'batches (in this case, %d batches).' %
steps_per_epoch * epochs)
break
batch_logs.update(outputs)
callbacks.on_batch_end(step_index, batch_logs)
step_index = step_index + step_count
if callbacks.model.stop_training:
break
if (do_validation and
training_utils.should_run_validation(validation_freq, epoch)):
logging.info('Running validation at fit epoch: %s', epoch)
if model._compile_distribution:
# Since we create a new clone from the original model we need to copy
# the weights back to the original model before we can run validation.
distributed_training_utils._copy_weights_to_original_model(
model, model._distributed_model_train, ModeKeys.TRAIN)
val_outs = experimental_tpu_test_loop( # pylint: disable=undefined-variable
model,
val_dataset,
steps=validation_steps,
verbose=verbose)
if not isinstance(val_outs, list):
val_outs = [val_outs]
# Same labels assumed.
for label, val_out in zip(out_labels, val_outs):
epoch_logs['val_' + label] = val_out
callbacks.on_epoch_end(epoch, epoch_logs)
if callbacks.model.stop_training:
break
callbacks.on_train_end()
if model._compile_distribution:
# Copy the weights back from the replicated model to the original model.
distributed_training_utils._copy_weights_to_original_model(
model, model._distributed_model_train, ModeKeys.TRAIN)
scope.__exit__(None, None, None)
return model.history
