def lr_finder(
model: tf.keras.Model,
optimizer: tf.keras.optimizers.Optimizer,
loss_fn: tf.keras.losses.Loss,
dataset,
learn_rates: LrGenerator,
losses: SmoothedLoss,
) -> Lr:
# To run the lr finder directly before training, we need to reset
# the initial weights. Here, I do it with building, storing and resetting.
model.build(dataset.element_spec[0].shape)
weights = model.get_weights()
for lr, (source, target) in zip(learn_rates(), dataset):
tf.keras.backend.set_value(optimizer.lr, lr)
loss = train_step(model, optimizer, loss_fn, source, target).numpy()
if losses.no_progress:
break
losses.update(loss)
model.set_weights(weights)
return Lr(learn_rates, losses)
def train(train_x_y, model:tf.keras.Model, val_xy=None, epochs=10, input_dim=None):
'''
TF高阶API
:param train_x_y: 训练集
:param model: 模型实例
:param epochs: 迭代次数
:return:
'''
if input_dim:
model.build(input_shape=(None, input_dim))
print(model.summary())
# 为训练选择优化器和损失函数
model.compile(optimizer=tf.keras.optimizers.Adam(lr=1e-3),
loss=tf.losses.BinaryCrossentropy(from_logits=False),
metrics=['binary_accuracy'])
# 训练
# model.fit(train_x_y, epochs=epochs, validation_data=val_xy, validation_freq=5)
model.fit(train_x_y, epochs=epochs)
return model
def add_spectral_norm_for_model(model: tf.keras.Model):
if model.built:
raise ValueError("Can't add spectral norm on built layer!")
original_build = model.build
# Very Very Evil HACK
def new_build(self, input_shape):
original_build(input_shape)
for sub_layer in self.layers:
add_spectral_norm(sub_layer)
model.build = types.MethodType(new_build, model)
def export_saved_model(model: tf.keras.Model,
input_shape: Tuple[int, int, int, int, int],
export_path: str = '/tmp/movinet/',
causal: bool = False,
bundle_input_init_states_fn: bool = True,
checkpoint_path: Optional[str] = None) -> None:
"""Exports a MoViNet model to a saved model.
Args:
model: the tf.keras.Model to export.
input_shape: The 5D spatiotemporal input shape of size
[batch_size, num_frames, image_height, image_width, num_channels].
Set the field or a shape position in the field to None for dynamic input.
export_path: Export path to save the saved_model file.
causal: Run the model in causal mode.
bundle_input_init_states_fn: Add init_states as a function signature to the
saved model. This is not necessary if the input shape is static (e.g.,
for TF Lite).
checkpoint_path: Checkpoint path to load. Leave blank to keep the model's
initialization.
"""
# Use dimensions of 1 except the channels to export faster,
# since we only really need the last dimension to build and get the output
# states. These dimensions can be set to `None` once the model is built.
input_shape_concrete = [1 if s is None else s for s in input_shape]
model.build(input_shape_concrete)
# Compile model to generate some internal Keras variables.
model.compile()
if checkpoint_path:
checkpoint = tf.train.Checkpoint(model=model)
status = checkpoint.restore(checkpoint_path)
status.assert_existing_objects_matched()
if causal:
# Call the model once to get the output states. Call again with `states`
# input to ensure that the inputs with the `states` argument is built
# with the full output state shapes.
input_image = tf.ones(input_shape_concrete)
_, states = model({
**model.init_states(input_shape_concrete), 'image':
input_image
})
_ = model({**states, 'image': input_image})
# Create a function to explicitly set the names of the outputs
def predict(inputs):
outputs, states = model(inputs)
return {**states, 'logits': outputs}
specs = {
name: tf.TensorSpec(spec.shape, name=name, dtype=spec.dtype)
for name, spec in model.initial_state_specs(input_shape).items()
}
specs['image'] = tf.TensorSpec(input_shape,
dtype=model.dtype,
name='image')
predict_fn = tf.function(predict, jit_compile=True)
predict_fn = predict_fn.get_concrete_function(specs)
init_states_fn = tf.function(model.init_states, jit_compile=True)
init_states_fn = init_states_fn.get_concrete_function(
tf.TensorSpec([5], dtype=tf.int32))
if bundle_input_init_states_fn:
signatures = {'call': predict_fn, 'init_states': init_states_fn}
else:
signatures = predict_fn
tf.keras.models.save_model(model, export_path, signatures=signatures)
else:
_ = model(tf.ones(input_shape_concrete))
tf.keras.models.save_model(model, export_path)
