def load_weight(self, base_directory: str, filename: str,
model: keras.models.Model) -> keras.models.Model:
if filename == '':
model.load_weights(
os.path.join(base_directory, self.weight_filename))
else:
model.load_weights(filename)
return model
def pretrain_logger(args, model: keras.models.Model, train_sequence, test_sequence):
train_metrics = dict(
zip(model.metrics_names, model.evaluate(train_sequence)))
val_metrics = dict(zip(
("val_" + name for name in model.metrics_names), model.evaluate(test_sequence)))
checkpoint_path = args.checkpoint_dir.format(
epoch=0, **train_metrics, **val_metrics)
print("Pretrain: saving model to {}".format(checkpoint_path))
keras.models.save_model(model, checkpoint_path)
def save_model(model: keras.models.Model, topology: dict,
filepath: str) -> None:
"""
Save a model to a file (tf.keras models only)
The method save the model topology, as given as a
Args:
model: model object
topology (dict): a dictionary of topology elements and their values
filepath (str): path to save model
"""
with tempfile.NamedTemporaryFile(suffix='.h5', delete=True) as fd:
model.save_weights(fd.name)
model_weights = fd.read()
data = {'model_weights': model_weights, 'model_topology': topology}
with open(filepath, 'wb') as fp:
pickle.dump(data, fp)
def plot_image_and_reconstruction(X: np.ndarray,
model: keras.models.Model,
idx: Optional[int] = None,
cmap: Optional[str] = 'binary') -> None:
"""Plot an image and its reconstruction from an autoencoder model
Args:
X (np.ndarray): Training data of autoencoder
model (keras.models.Model): Autoencoder
idx (Optional[int], optional): The index of the training sample to plot. Defaults to None, in which case a random index is chosen.
cmap (Optional[str], optional): The image mapping for pyplot. Defaults to 'binary'.
"""
idx = idx or np.random.choice(len(X))
image = X[idx]
[reconstruction] = model.predict(image[np.newaxis, :])
difference = np.abs(image - reconstruction)
fig = plt.figure(figsize=(3 * 1.5, 3))
plt.subplot(1, 3, 1)
plt.imshow(image, cmap=cmap)
plt.subplot(1, 3, 2)
plt.imshow(reconstruction, cmap=cmap)
plt.subplot(1, 3, 3)
plt.imshow(difference, cmap=cmap)
plt.show()
def _hamiltonian(
self, model: keras.models.Model,
inputs: tf.Tensor, targets: tf.Tensor,
params: tf.Tensor, momenta: tf.Tensor
) -> float:
"""
Calculate the value of the Hamiltonian of the system.
The parameter state should be a tf.Tensor
(x in the thesis, ξ in the original paper).
It takes two arguments, params and momenta, instead of the single state
argument, since it is possible to faster calculate the partial derivative
with regard to the momenta.
"""
# Assign state parameters to model
_model_set_flat_variables(model, params)
# Predict the outcome with the new parameters
prediction = model(inputs)
return 0.5 * (
self.alpha * model.loss(targets, prediction)
+ self.beta * tf.tensordot(params, params, 2)
+ self.gamma * tf.tensordot(momenta, momenta, 2)
)
def wrap(self, model: keras.models.Model):
with K.name_scope("training"):
with K.name_scope("Lookahead"):
# initialize counter and slow_weights
self.count = tf.Variable(0, dtype=tf.int32, trainable=False, name="update_count")
K.track_variable(self.count)
self.slow_weights = []
for fast_param in model.trainable_weights:
with ops.control_dependencies([fast_param]):
slow_param = tf.Variable(fast_param.initialized_value(),
dtype=fast_param.dtype,
trainable=False,
name=fast_param.name.split(":")[0])
self.slow_weights.append(slow_param)
K.track_variable(slow_param)
def lookahead_update():
with K.name_scope("training"):
with K.name_scope("Lookahead"):
# count++ mod k
count_op = state_ops.assign_add(self.count, 1)
def fast_update():
return control_flow_ops.no_op()
def slow_update():
with ops.control_dependencies(model.trainable_weights):
slow_ups = [state_ops.assign_add(slow, (fast - slow) * self.alpha)
for slow, fast in zip(self.slow_weights, model.trainable_weights)]
with ops.control_dependencies(slow_ups):
fast_ups = [state_ops.assign(fast, slow)
for slow, fast in zip(self.slow_weights, model.trainable_weights)]
return control_flow_ops.group(*fast_ups)
with ops.control_dependencies([count_op]):
update_op = control_flow_ops.cond(math_ops.equal(math_ops.mod(self.count, self.k), 0),
slow_update, fast_update)
return update_op
model.add_update(lookahead_update)
return model
def split_dataset(dataset: Dataset, model: keras.models.Model) -> Split:
dataset_gen = data_flow_from_disk(SOURCE_PATH, dataset, label_indices,
False, BS, SEED, MODEL)
predictions: Predictions = OrderedDict()
for dataset_item, prediction in zip(dataset.keys(),
model.predict(dataset_gen)):
predictions[dataset_item] = prediction
return SPLITTER_FACTORY(predictions, 50)(dataset)
def save_model_tfjs(model: ks.models.Model, save_loc: str) -> None:
word_pred_model = ks.models.Model(
inputs=model.inputs,
outputs=[model.get_layer('word_preds').output],
)
word_pred_model.compile(optimizer='sgd', loss='mse')
tfjs.converters.save_keras_model(
word_pred_model,
save_loc,
)
def get_excitations(model: ks.models.Model, num_sentences: int,
dataset: Dataset, save_loc: str) -> None:
conv_len = model.get_layer('convs').get_weights()[0].shape[0]
conv_output_model = ks.models.Model(
inputs=model.inputs,
outputs=[model.get_layer('convs').output],
)
preds = conv_output_model.predict(dataset.x_train[:num_sentences])
def get_preds(idxs: Matrix, preds: Matrix):
num_padding = np.sum(idxs == 0)
idxs, preds = idxs[num_padding:], preds[num_padding + 1:]
words = dataset.decode(idxs)
return {
'words': words,
'preds': [[float(j) for j in i] for i in preds.T],
}
excitations = {
'conv_len':
conv_len,
'activations': [
get_preds(i, p)
for i, p in zip(dataset.x_train[:num_sentences], preds)
],
}
utils.save_json(excitations, save_loc, 'excitations.json')
with open(os.path.join(save_loc, 'excitations_show.txt'), 'w') as f:
for example in excitations['activations']:
for i in range(3):
ordered_sentences = get_ordered_sentences(
example['words'],
example['preds'][i],
conv_len,
)
f.write('Layer {}\n'.format(i))
f.write('\n'.join('{}: {}'.format(*g)
for g in ordered_sentences))
f.write('\n\n')
def load_checkpoint(self, fs: FSBase, model: keras.models.Model) -> None:
with tempfile.NamedTemporaryFile(suffix=".h5") as tf:
local_fs = FileSystem()
with fs.open("model.h5", "rb") as fin:
local_fs.writefile(tf.name, fin)
model.load_weights(tf.name)
def save_checkpoint(self, fs: FSBase, model: keras.models.Model) -> None:
with tempfile.NamedTemporaryFile(suffix=".h5") as tf:
model.save_weights(tf.name)
with open(tf.name, "rb") as fin:
fs.writefile("model.h5", fin)
def save_model(self, base_directory: str, model: keras.models.Model):
open(os.path.join(base_directory, self.model_config_filename),
'w').write(model.to_yaml())
def save_weight(self, base_directory: str, model: keras.models.Model):
print('saving weights...')
model.save_weights(os.path.join(base_directory, self.weight_filename))
def save_model(model_name: str, model: keras.models.Model):
model_json = model.to_json()
with open(model_name, "w") as json_file:
json_file.write(model_json)
def predict_future_vals(model: keras.models.Model, window_frame: np.array):
future_window = model.predict(window_frame)
future_window_rounded = np.rint(future_window)
future_window_rounded = np.squeeze(future_window_rounded)
return future_window_rounded
def train_batch(
self, model: keras.models.Model,
input_batch: tf.Tensor, target_batch: tf.Tensor,
metrics=[]
) -> float:
"""
Train the model with a single batch of samples
Parameters
----------
model : tf.keras.Model
The model to train
input_batch : tf.Tensor
Input samples of the batch to train.
These contain the inputs of the batch
target_batch : tf.Tensor
Output samples of the batch to train.
These contain the expected output of the model for the inputs.
metrics : List[tf.metrics.Metric]
List of tensorflow metrics.
These will be evaluated after the training.
"""
# Basic sanity check that input sample count matches target sample
# count
sample_cnt = input_batch.shape[0]
assert sample_cnt == target_batch.shape[0]
# Call the model handler
# (no-op if the model didn't change since the last iteration)
self._check_model_and_state(model)
# Redefine hamiltonian and gradient for the current batch
batch_hamiltonian = self.get_hamiltonian(model, input_batch, target_batch)
@tf.function
def loss_gradient(params: tf.Tensor) -> tf.Tensor:
_model_set_flat_variables(model, params)
with tf.GradientTape() as tape:
tape.watch(input_batch)
loss = 0.5 * (
self.alpha * model.loss(target_batch, model(input_batch, training=True)) \
+ self.beta * tf.tensordot(params, params, 2)
)
return _flatten_variables(
tape.gradient(loss, model.trainable_variables)
)
self._integrator.loss_gradient = loss_gradient
params = _flatten_variables(model.trainable_variables)
params, velocity = self._integrator.integrate(
self.ivp_period, self.ivp_step_size, params
)
momenta = tf.sparse.sparse_dense_matmul(self.M, velocity)
_model_set_flat_variables(model, params)
# Predict output for the current batch and update the metrics accordingly
prediction = model(input_batch)
for metric in metrics:
metric.update_state(target_batch, prediction)
# We return the loss for the current batch and the energy in the system
# for the current batch
return model.loss(target_batch, prediction), batch_hamiltonian(params, momenta)