def load_bert_weight_from_ckpt(
bert_model: tf.keras.Model,
bert_ckpt_dir: str,
repl_patterns: Optional[Dict[str, str]] = None
) -> Tuple[tf.keras.Model, Mapping[str, str], Mapping[str, Any]]:
"""Loads checkpoint weights and match to model weights if applicable.
Args:
bert_model: The BERT Encoder model whose weights to load from checkpoints.
bert_ckpt_dir: Path to BERT pre-trained checkpoints.
repl_patterns: A mapping of regex string patterns and their replacements. To
be used to update checkpoint weight names so they match those in
bert_model (e.g., via re.sub(pattern, repl, weight_name))
Returns:
bert_model: The BERT Encoder model with loaded weights.
names_to_keys: A dict mapping of weight name to checkpoint keys.
keys_to_weights: A dict mapping of checkpoint keys to weight values.
"""
# Load a dict mapping of weight names to their corresponding checkpoint keys.
names_to_keys = object_graph_key_mapping(bert_ckpt_dir)
if repl_patterns:
# Update weight names so they match those in bert_model
names_to_keys = {
update_weight_name(repl_patterns, weight_name): weight_key
for weight_name, weight_key in names_to_keys.items()
}
# Load a dict mapping of checkpoint keys to weight values.
logging.info('Loading weights from checkpoint: %s', bert_ckpt_dir)
keys_to_weights = load_ckpt_keys_to_weight_mapping(bert_ckpt_dir)
# Arranges the pre-trained weights in the order of model weights.
init_weight_list = match_ckpt_weights_to_model(bert_model, names_to_keys,
keys_to_weights)
# Load weights into model.
bert_model.set_weights(init_weight_list)
return bert_model, names_to_keys, keys_to_weights
def evaluate_neural_network(testset: tf.data.Dataset,
neural_network: tf.keras.Model,
num_monte_carlo: int = 0):
assert isinstance(num_monte_carlo, int)
assert num_monte_carlo >= 0
if num_monte_carlo > 0:
assert num_monte_carlo > 1
y_true = []
probabilities = []
for batch_X, batch_y in testset:
assert batch_y.shape[0] == batch_X[0].shape[0]
y_true.append(batch_y)
if num_monte_carlo > 0:
new_probabilities = tf.reduce_mean(tf.stack([
neural_network.predict_on_batch(batch_X)
for _ in range(num_monte_carlo)
]),
axis=0)
else:
new_probabilities = neural_network.predict_on_batch(batch_X)
probabilities.append(
np.reshape(new_probabilities.numpy(),
newshape=(batch_X[0].shape[0], )))
y_true = np.concatenate(y_true)
probabilities = np.concatenate(probabilities)
print('Evaluation results:')
print(' ROC-AUC is {0:.6f}'.format(roc_auc_score(y_true,
probabilities)))
y_pred = np.asarray(probabilities >= 0.5, dtype=np.uint8)
del probabilities
with warnings.catch_warnings():
warnings.filterwarnings("ignore", category=UndefinedMetricWarning)
print(' Precision is {0:.6f}'.format(
precision_score(y_true, y_pred, average='binary')))
print(' Recall is {0:.6f}'.format(
recall_score(y_true, y_pred, average='binary')))
print(' F1 is {0:.6f}'.format(
f1_score(y_true, y_pred, average='binary')))
print('')
del y_pred, y_true
def add_image_loss(
model: tf.keras.Model,
fixed_image: tf.Tensor,
pred_fixed_image: tf.Tensor,
loss_config: dict,
) -> tf.keras.Model:
"""
add image dissimilarity loss of ddf into model
:param model: tf.keras.Model
:param fixed_image: tensor of shape (batch, f_dim1, f_dim2, f_dim3)
:param pred_fixed_image: tensor of shape (batch, f_dim1, f_dim2, f_dim3)
:param loss_config: config for loss
"""
if loss_config["dissimilarity"]["image"]["weight"] > 0:
loss_image = tf.reduce_mean(
image_loss.dissimilarity_fn(
y_true=fixed_image,
y_pred=pred_fixed_image,
**loss_config["dissimilarity"]["image"],
))
weighted_loss_image = (loss_image *
loss_config["dissimilarity"]["image"]["weight"])
model.add_loss(weighted_loss_image)
model.add_metric(loss_image,
name="loss/image_dissimilarity",
aggregation="mean")
model.add_metric(
weighted_loss_image,
name="loss/weighted_image_dissimilarity",
aggregation="mean",
)
return model
def hparam_search(
hparam_options: Dict,
model: tf.keras.Model,
dataset: DataSet,
log_root = None,
verbose = False,
debug = False
) -> List[Dict]:
def onexit(log_dir):
print('Ctrl-C KeyboardInterrupt')
if os.path.isdir(log_dir):
shutil.rmtree(log_dir) # remove logs for incomplete trainings
print(f'rm -rf {log_dir}')
model_config = model.get_config()
combninations = hparam_combninations(hparam_options)
logdir = hparam_logdir(combninations[0], hparam_options, log_root)
stats_history = []
# print(f"--- Model Config: ", model_config)
print(f"--- Testing {len(combninations)} combinations in {logdir}")
print(f"--- hparam_options: ", hparam_options)
for index, hparams in enumerate(combninations):
run_name = hparam_run_name(hparams, hparam_options)
logdir = hparam_logdir(hparams, hparam_options, log_root)
print("")
print(f"--- Starting trial {index+1}/{len(combninations)}: {logdir.split('/')[-2]} | {run_name}")
print(hparams)
if os.path.exists(logdir):
print('Exists: skipping')
continue
if debug: continue
atexit.register(onexit, logdir)
# DOCS: https://www.tensorflow.org/guide/keras/save_and_serialize
if model_config['name'] == 'sequential': model_clone = tf.keras.Sequential.from_config(model_config)
else: model_clone = tf.keras.Model.from_config(model_config)
stats = hparam.model_compile_fit(hparams, model_clone, dataset, log_dir=logdir, verbose=verbose)
stats_history += stats
print(stats)
atexit.unregister(onexit)
print("")
print("--- Stats History")
print(stats_history)
print("--- Finished")
return stats_history
def gen_text(transformer_model: tf.keras.Model,
tokenizer: AutoTokenizer,
config,
dataset,
input_text: str,
temperature=0.3,
topk=300):
input_token_dic = tokenizer(input_text,
truncation=True,
max_length=512,
return_tensors='tf')
if input_token_dic['input_ids'].shape[1] > config.MAX_LENGTH:
input_token_dic['input_ids'] = input_token_dic[
'input_ids'][:, -config.MAX_LENGTH:]
input_token_dic['attention_mask'] = input_token_dic[
'attention_mask'][:, -config.MAX_LENGTH:]
input_token_dic['token_type_ids'] = input_token_dic[
'token_type_ids'][:, -config.MAX_LENGTH:]
else:
input_token_dic = tokenizer(input_text,
padding='max_length',
truncation=True,
max_length=config.MAX_LENGTH,
return_tensors='tf')
output_ids = np.full((1, config.MAX_CHAR_LEN + 1),
dataset.char_idx[dataset.MU])
output_ids[0, 0] = dataset.char_idx[dataset.STR]
generated_text = ""
for cur in range(1, config.MAX_CHAR_LEN + 1):
# encoder_output = encoder_model.predict_on_batch((input_token_dic['input_ids'], input_token_dic['attention_mask']))
# preds = decoder_model.predict_on_batch(
# (encoder_output, output_ids[:, :-1]))
preds = transformer_model.predict_on_batch(
(input_token_dic['input_ids'], input_token_dic['attention_mask'],
output_ids[:, :-1]))
temp_ids = []
for pidx in range(config.MAX_CHAR_LEN):
temp_ids.append(
int(sample(preds[0][pidx], temperature=temperature,
topk=topk)))
next_id = temp_ids[cur - 1]
if next_id == dataset.char_idx[dataset.MU]:
break
output_ids[0, cur] = next_id
generated_text += dataset.idx_char[next_id]
return generated_text
def make_prediction(model: tf.keras.Model, input: str) -> np.ndarray:
if os.path.isdir(input):
raise IsADirectoryError()
# Loading data
image = Image.open(input)
data = np.array(image, dtype=np.float64)
# Expanding to add batch dimension
data = np.expand_dims(data, axis=0)
# Making a prediction and converting to uint8
prediction = (model.predict(data) * 255)[0]
return prediction.astype(np.uint8)
def load_weights(model: tf.keras.Model,
model_weights_path: str,
checkpoint_format: str = 'tf_checkpoint'):
"""Load model weights from the given file path.
Args:
model: the model to load weights into
model_weights_path: the path of the model weights
checkpoint_format: The source of checkpoint files. By default, we assume the
checkpoint is saved by tf.train.Checkpoint().save(). For legacy reasons,
we can also resotre checkpoint from keras model.save_weights() method by
setting checkpoint_format = 'keras_checkpoint'.
"""
if checkpoint_format == 'tf_checkpoint':
checkpoint_dict = {'model': model}
checkpoint = tf.train.Checkpoint(**checkpoint_dict)
checkpoint.restore(model_weights_path).assert_existing_objects_matched()
elif checkpoint_format == 'keras_checkpoint':
# Assert makes sure load is successeful.
model.load_weights(model_weights_path).assert_existing_objects_matched()
else:
raise ValueError(f'Unsupported checkpoint format {checkpoint_format}.')
def measure_backwards_fps(model: tf.keras.Model,
data: Tuple[Tuple[tf.Tensor]],
repeats: int = 100,
verbose: int = 1):
times = []
for i in range(repeats + 1):
start = time.time()
model.train_step(data)
times.append(time.time() - start)
if verbose:
print(f"\rProgress: {(i + 1) / 101:>7.2%}", end="")
if verbose:
print()
mean_time = tf.reduce_mean(times[1:]).numpy()
if verbose:
print(f" [*] BATCH PER SEC: {1 / mean_time:.5f}")
print(f" [*] SEC PER BATCH: {mean_time:.5f}")
print()
return mean_time
def custom_recall_validation_with_generator(
self, generator: Iterable[Tuple[np.ndarray, np.ndarray]],
model: tf.keras.Model) -> float:
# TODO: write description
total_predictions = np.zeros((0, ))
total_ground_truth = np.zeros((0, ))
for x, y in generator:
predictions = model.predict(x)
predictions = predictions.argmax(axis=-1).reshape((-1, ))
total_predictions = np.append(total_predictions, predictions)
total_ground_truth = np.append(total_ground_truth,
y.argmax(axis=-1).reshape((-1, )))
return self.evaluation_metric(total_ground_truth, total_predictions)
def set_model(self, model: tf.keras.Model):
# pylint: disable=protected-access
old_model = getattr(self, "model", None)
if old_model is not None:
del old_model._callbacks[self.name]
if not hasattr(model, "_callbacks"):
model._callbacks = {}
callbacks = model._callbacks
# pylint: enable=protected-access
assert self.name not in callbacks
callbacks[self.name] = self
self.model = model # pylint: disable=attribute-defined-outside-init
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 _predict_img(self, model: tf.keras.Model, prediction_img_amount: int):
pairs = {'frames': [], 'masks': []}
for _ in range(prediction_img_amount):
sample_pair = next(self.DsGenerator.get_next_pair())
pairs['frames'].append(sample_pair['frame'])
pairs['masks'].append(sample_pair['mask'])
feature_maps = model.predict([pairs['frames']])
for i, feature_map in enumerate(feature_maps):
predicted_mask = create_mask(feature_map)
self.draw_prediction(pairs['frames'][i],
np.reshape(pairs['masks'][i], pairs['masks'][i].shape[:-1]),
predicted_mask, i)
def run_time_priorized_detection(time_data_loader: COCODoomDataset,
model: tf.keras.Model,
detection_file="default"):
def process(file_names_list):
file_name_dataset = tf.data.Dataset.from_tensor_slices(file_names_list)
file_name_dataset = file_name_dataset.map(tf.io.read_file)
file_name_dataset = file_name_dataset.map(tf.io.decode_image)
file_name_dataset = file_name_dataset.map(
lambda t: tf.image.convert_image_dtype(t, tf.float32))
return file_name_dataset
img_root = time_data_loader.cfg.images_root
file_names = [
os.path.join(img_root, meta["file_name"])
for meta in time_data_loader.image_meta.values()
if "prev_image_id" in meta
]
category_index = {
cat["name"]: cat
for cat in time_data_loader.categories.values()
}
prev_file_names = []
for meta in time_data_loader.image_meta.values():
prev_meta_id = meta.get("prev_image_id", None)
if prev_meta_id is None:
continue
prev_meta = time_data_loader.image_meta[prev_meta_id]
prev_file_names.append(os.path.join(img_root, prev_meta["file_name"]))
present: tf.data.Dataset = process(file_names)
past: tf.data.Dataset = process(prev_file_names)
IDs = tf.data.Dataset.from_tensor_slices(
[meta["id"] for meta in time_data_loader.image_meta.values()])
dataset = tf.data.Dataset.zip((present, past, IDs))
detections = []
print()
for i, (past, present, ID) in enumerate(dataset, start=1):
result = model.detect((past, present))
detections.extend(
_convert_to_coco_format(result, category_index, int(ID)))
print("\r [Verres] - COCO eval "
f"P: {i / len(time_data_loader.image_meta):>7.2%} ")
print()
return _evaluate(detections, time_data_loader, detection_file)
def save_tf_model(model: tf.keras.Model,
filepath: str,
save_dir: str = None,
model_name: str = 'model') -> None:
"""
Save TensorFlow model.
Parameters
----------
model
A tf.keras Model.
filepath
Save directory.
save_dir
Save folder.
model_name
Name of saved model.
"""
# create folder for model weights
if not os.path.isdir(filepath):
logger.warning(
'Directory {} does not exist and is now created.'.format(filepath))
os.mkdir(filepath)
if not save_dir:
model_dir = os.path.join(filepath, 'model')
else:
model_dir = os.path.join(filepath, save_dir)
if not os.path.isdir(model_dir):
os.mkdir(model_dir)
# save classification model
if isinstance(model, tf.keras.Model) or isinstance(model,
tf.keras.Sequential):
model.save(os.path.join(model_dir, model_name + '.h5'))
else:
logger.warning(
'No `tf.keras.Model` or `tf.keras.Sequential` detected. No model saved.'
)
def transfer_weights(source_model: tf.keras.Model,
target_model: tf.keras.Model,
is_cloned: bool = False,
layer_name_prefix: str = '',
beta: float = 0.0) -> None:
"""Linear beta-interpolation of weights from source_model to target_model.
Can be used to maintain a shadow exponential moving average of source_model. Only weights of layers with the same
name in both models and both starting with 'layer_name_prefix' are transferred.
If target_model and source_model are clones and share the exact same topology a significantly faster implementation
is used. If is_cloned is False, this function assumes source_model is a topological sub-network of target_model; in
that case missing layers in either target_model or source_model are silently ignored.
Args:
source_model: the source to transfer weights from
target_model: the target to transfer weights to
is_cloned: whether or not source and target are exact clones (significantly speeds up computation)
layer_name_prefix: only layers, which names start with layer_name_prefix, are transferred
beta: value for linear interpolation; must be within [0.0, 1.0)
Raises:
ValueError: if beta exceeds interval [0.0, 1.0)
"""
if not 0.0 <= beta < 1.0:
raise ValueError(
f'beta must be within [0.0, 1.0) but received beta={beta} instead')
if is_cloned: # same exact layer order and topology in both models
for source_layer, target_layer in zip(source_model.layers,
target_model.layers):
if source_layer.name == target_layer.name and source_layer.name.startswith(
layer_name_prefix):
for source_var, target_var in zip(source_layer.variables,
target_layer.variables):
delta_value = (1 - beta) * (target_var - source_var)
target_var.assign_sub(delta_value)
else: # iterate source_model.layers and transfer to target_layer, if target_layer exists
for source_layer in source_model.layers:
source_vars = source_layer.variables
if source_layer.name.startswith(layer_name_prefix) and source_vars:
try:
target_layer = target_model.get_layer(
name=source_layer.name)
except ValueError:
continue
for source_var, target_var in zip(source_vars,
target_layer.variables):
delta_value = (1 - beta) * (target_var - source_var)
target_var.assign_sub(delta_value)
def add_compile(model: tf.keras.Model, **params) -> tf.keras.Model:
"""Compiles the model with the given parameters.
Parameters
----------
model: tf.keras.Model
The model to compile.
Returns
-------
tf.keras.Model
Compiles the model with loss and optimizer.
"""
optimizer = Adam(lr=params["learning_rate"],
clipnorm=params.get("clipnorm", 1))
model.compile(
loss=tf.keras.losses.BinaryCrossentropy(),
optimizer=optimizer,
metrics=["accuracy",
tf.keras.metrics.AUC(multi_label=True)],
)
def test_single_image(model: tf.keras.Model, path: str, save_result: bool):
"""
Loads image from disk and inference on it
:param model: Model object with weights loaded
:param path: Image path to load
:param save_result: Specifies if result will be save as image or not
"""
# Load image with PIL as RGB
print('Loading image!')
img_pil = Image.open(path).convert("RGB")
# Reshape for model input
img_pil_resized = img_pil.resize((224, 224))
# Normalize
img = np.array(img_pil_resized, dtype=np.float32) / 255.
# Creating empty axis on batch axis
img = img[np.newaxis, :, :, :]
# Inference on model
t = time.time()
print('Inference started')
pred = model.predict(img)
print(f'Inference finished: {time.time() - t:.5f} second\n\n')
# Unpacking predictions
is_defect_pred = pred[0]
class_type_pred = pred[1]
bbox_param_pred = np.squeeze(pred[2])
bbox_center_pred = np.squeeze(pred[3])
# Printing logs
print(f'Is defected: \tPrediction | {is_defect_pred.squeeze()}')
print(
f'Class type: \tPrediction | {np.argmax(class_type_pred.squeeze()) + 1}'
)
print(f'Bbox center: \tPrediction | {bbox_center_pred}')
print(f'Bbox params: \tPrediction | {bbox_param_pred}')
img_original = np.array(img_pil, dtype=np.uint8)
# Plot configuration
fig, ax = plt.subplots(figsize=(12, 12))
plt.imshow(img_original, cmap='gray')
# Creating un-filled ellipse on image
e = Ellipse(xy=(bbox_center_pred * 512),
width=bbox_param_pred[0] * 256,
height=bbox_param_pred[1] * 256,
angle=((bbox_param_pred[2] * 2 * np.pi - np.pi) * 180 / np.pi),
edgecolor='b',
lw=2,
facecolor='none')
e.set_alpha(0.8)
ax.add_artist(e)
if save_result:
# Saves plot to disk
fig.savefig('logs/result.png')
def attack_by_bit_success(
x: np.ndarray, img: np.ndarray, y_true: int, model: tf.keras.Model
):
"""Predict ONE image and return True if expected. None otherwise."""
attack_image = original_perturb_image_by_bit_fault(x, img)
confidence = model.predict(attack_image)[0]
predicted_class = np.argmax(confidence)
# If the prediction is what we want (misclassification or
# targeted classification), return True
logging.debug(f"Confidence: {confidence[y_true]}")
if predicted_class == y_true:
return True
def train_model(model: tf.keras.Model,
train_ds: tf.data.Dataset,
val_ds: tf.data.Dataset) -> tf.keras.Model:
"""Function trains the model and evaluates it's performance,
displays training metrics.
:param model: Untrained model
:param train_ds: Training Dataset containing input data and labels
:param val_ds: Validation Dataset containing input data and labels
:return: Trained model
"""
early_stop = tf.keras.callbacks.EarlyStopping(monitor='val_loss',
patience=5,
restore_best_weights=True)
history = model.fit(train_ds, validation_data=val_ds,
epochs=100, verbose=2, callbacks=[early_stop],
use_multiprocessing=True, workers=-1)
loss, acc = model.evaluate(valid_ds)
print(f'Validation loss: {loss}\nValidation accuracy: {acc}')
plot_history(history)
return model
def get_logits_labels(model: tf.keras.Model,
evaluation_set: tf.data.TFRecordDataset):
"""Predicts model output on the given evaluation set
Parameters:
model: tf.keras.Model to be used for prediction
evaluation_set: tf.data.TFRecordDataset to be used for prediction
Returns:
model output: numpy.ndarray
labels: numpy.ndarray
"""
logits_numpys = model.predict(evaluation_set).squeeze()
labels_numpys = np.concatenate(list(map(
lambda x: x[1], evaluation_set))).squeeze().argmax(axis=-1)
return logits_numpys, labels_numpys
def train_and_eval(model: tf.keras.Model, model_dir: str,
train_input_dataset: tf.data.Dataset,
eval_input_dataset: tf.data.Dataset, steps_per_epoch: int,
epochs: int, eval_steps: int):
"""Train and evaluate."""
callbacks = get_callbacks(model, model_dir)
history = model.fit(x=train_input_dataset,
validation_data=eval_input_dataset,
steps_per_epoch=steps_per_epoch,
epochs=epochs,
validation_steps=eval_steps,
callbacks=callbacks)
tf.get_logger().info(history)
return model
def train_model(model: tf.keras.Model, dataset: pd.DataFrame, epochs: int,
batch_size: int,
target_name: str) -> tf.keras.callbacks.History:
"""Trains the given keras model with the given dataset."""
features = {name: np.array(value) for name, value in dataset.items()}
target = np.array(features.pop(target_name))
history = model.fit(x=features,
y=target,
batch_size=batch_size,
epochs=epochs,
validation_split=0.2,
verbose=1)
return history
def freeze_all(model: tf.keras.Model, frozen: bool = True):
"""
freeze_all(model: tf.Keras.Model, frozen: boolean)
A function which takes tf Model and makes specified layers non trainable for transfer learning
:param model: Trained tensorflow model
:param frozen: boolean to freeze the model layers
:return: returns void with the specified model layers are made non trainable
"""
model.trainable = not frozen
if isinstance(model, tf.keras.Model):
for l in model.layers:
freeze_all(l, frozen)
def train_model(
model: tf.keras.Model,
train_dataset: tf.data.Dataset,
hyperparameters: Hyperparameters,
) -> Tuple[tf.keras.Model, str]:
# define the checkpoint directory to store the checkpoints
checkpoint_dir = "./training_checkpoints"
# define the name of the checkpoint files
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt_{epoch}")
# define a callback for printing the learning rate at the end of each epoch
class PrintLR(tf.keras.callbacks.Callback):
def on_epoch_end(self, epoch, logs=None):
print(
"\nLearning rate for epoch {} is {}".format(
epoch + 1, model.optimizer.lr.numpy()
)
)
# put all the callbacks together
callbacks = [
tf.keras.callbacks.TensorBoard(log_dir="./logs"),
tf.keras.callbacks.ModelCheckpoint(
filepath=checkpoint_prefix, save_weights_only=True
),
tf.keras.callbacks.LearningRateScheduler(decay),
PrintLR(),
]
# train the model
model.fit(train_dataset, epochs=hyperparameters.epochs, callbacks=callbacks)
# save the model
model.save(MODEL_FILE_PATH, save_format="tf")
return model, checkpoint_dir
def build_checkpoint_callback(
model: tf.keras.Model,
dataset: tf.data.Dataset,
log_dir: str,
save_period: int,
ckpt_path: str,
) -> Tuple[CheckpointManagerCallback, int]:
"""
Function to prepare callbacks for training.
:param model: model to train
:param dataset: dataset for training
:param log_dir: directory of logs
:param save_period: save the checkpoint every X epochs
:param ckpt_path: path to restore ckpt
:return: a list of callbacks
"""
# fit the model for 1 step to initialise optimiser arguments as trackable Variables
model.fit(
x=dataset,
steps_per_epoch=1,
epochs=1,
verbose=0,
)
checkpoint_manager_callback = CheckpointManagerCallback(model,
log_dir + "/save",
period=save_period)
if ckpt_path:
initial_epoch_str = ckpt_path.split("-")[-1]
assert initial_epoch_str.isdigit(), (
f"Checkpoint path for checkpoint manager "
f"must be of form ckpt-epoch_count, got {ckpt_path}")
initial_epoch = int(initial_epoch_str)
checkpoint_manager_callback.restore(ckpt_path)
else:
initial_epoch = 0
return checkpoint_manager_callback, initial_epoch
def compare_generated_to_real(
dataloader,
num_images: int,
noise_size: int,
model: tf.keras.Model,
save_location: str,
img_size: int,
subsequent_model: Optional[tf.keras.Model] = None,
):
""" For a given number of images, generate the stackGAN stage 1 output by randomly sampling a dataloader.
The generated images and the real original are saved side-by-side in the save_location.
"""
rmdir(save_location)
mkdir(save_location)
noise_list = [
np.random.normal(0, 1, (1, noise_size)).astype("float32")
for idx in range(num_images)
]
samples = sample_data(dataloader,
num_samples=num_images,
img_size=img_size)
real_tensors, real_embeddings = zip(*samples)
stage1_tensors = [
model.generator([embedding, noise], training=False)[0]
for embedding, noise in zip(real_embeddings, noise_list)
]
real_images = format_as_images(real_tensors, is_real=True)
stage1_images = format_as_images(stage1_tensors, is_real=False)
if subsequent_model is not None:
stage2_tensors = [
subsequent_model.generator([generated_image, embedding],
training=False)[0] for generated_image,
embedding in zip(stage1_tensors, real_embeddings)
]
stage2_images = format_as_images(stage2_tensors, is_real=False)
for i, (real_image, stage1_image, stage2_image) in enumerate(
zip(real_images, stage1_images, stage2_images)):
image = concate_horizontallly(real_image,
stage1_img=stage1_image,
stage2_img=stage2_image)
image.save(os.path.join(save_location, f"fake-vs-real-{i}.png"))
else:
for i, (real_image,
stage1_image) in enumerate(zip(real_images, stage1_images)):
image = concate_horizontallly(real_image, stage1_img=stage1_image)
image.save(os.path.join(save_location, f"fake-vs-real-{i}.png"))
def export_best_model(strategy, pmetric, best_pmetric_val,
model: tf.keras.Model, out_dir):
"""Exports best model if current primary metric value is better than the best metric value"""
pmetric_name = pmetric.name
pmetric_val = float_metric_value(pmetric)
if best_pmetric_val is not None:
logging.info(
f'Existing best primary metric ({pmetric_name}): {best_pmetric_val}'
)
logging.info(f'Current primary metric ({pmetric_name}): {pmetric_val}')
if best_pmetric_val is None or pmetric_val > best_pmetric_val:
logging.info(
f'Exporting model with new best primary metric ({pmetric.name}): {pmetric_val}'
)
if should_export_checkpoint(strategy):
model.save(get_export_dir(out_dir))
else:
# In multi worker training we need every worker to save models, because variables can trigger synchronization on read and synchronization needs
# all workers to participate. To avoid workers overriding each other we save to a temporary directory on non-chief workers.
tmp_dir = tempfile.mkdtemp()
model.save(get_export_dir(out_dir))
tf.io.gfile.rmtree(tmp_dir)
return pmetric_val
return best_pmetric_val
def train(model: tf.keras.Model, model_path: str, n_epoch: int = 10):
x_train, x_test, y_train, y_test = load_and_split()
optimizer = tf.keras.optimizers.Adam(learning_rate=0.005)
model.compile(optimizer=optimizer,
loss=simnet_loss)
reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.2,
verbose=1, patience=5, min_lr=0.0001)
mcp_save = ModelCheckpoint(model_path, # {epoch:02d}-{val_loss:.2f}.hdf5
save_best_only=True, monitor='val_loss', mode='min')
# fit and check validation data
history = model.fit(x_train, y_train,
batch_size=2, epochs=n_epoch, workers=8,
callbacks=[reduce_lr, mcp_save],
validation_data=(x_test, y_test))
# model.save('./benchmarks/encoder_' + str(np.around(history.history['val_loss'][-1], 3)))
model.summary()
plot_loss(history)
tf.keras.utils.plot_model(model, 'simnet_model.png', show_shapes=True, expand_nested=True)
def train_model(
model:tf.keras.Model,
dataset: pd.DataFrame,
epochs: int,
batch_size: int,
label_name: str) -> tf.keras.callbacks.History:
"""Trains the given model on the dataset."""
features = {name: np.array(value) for name, value in dataset.items()}
label = np.array(features.pop(label_name))
history = model.fit(
x = features, y = label, epochs = epochs, batch_size = batch_size,
validation_split = 0.2, shuffle = True
)
return history
def read_dis_data(self, gen: tf.keras.Model):
print('generate_fake_data')
lens = len(self.x)
sizes = lens // 10240
y_fake = []
for i in range(sizes):
out = gen.predict(self.x[i * 10240: (i + 1) * 10240, :])
[y_fake.append([tf.argmax(o)]) for o in out]
print('\r%d/%d' % (i, sizes - 1), end='')
out = gen.predict(self.x[sizes * 10240:, :])
[y_fake.append([tf.argmax(o)]) for o in out]
true = []
fake = []
a = (self.x == 0).argmax(axis=1)
sizes = len(a)
for i, v in enumerate(a):
temp = self.x[i].copy()
temp[v] = self.y[i][0]
true.append(temp)
temp = self.x[i].copy()
temp[v] = y_fake[i][0]
fake.append(temp)
print('\r%d/%d' % (i, sizes - 1), end='')
return np.array(true), np.array(fake)