def train_model(model: tf.keras.Model, X_train, y_train,
epochs: int = 10,
tensorboard_path: str = './tensorboard/logs',
validation_data=None):
"""Function to fit the training data to the tf.keras model.
:param model: to train
:param X_train: values to train on
:param y_train: labels to train with
:param epochs: How many epochs the model should be trained
:param tensorboard_path: path to tensorboard
:param validation_data: validation data to pass to fit method
"""
# Single callback to get the Tensorboard logs
callbacks = []
tb_callback = tf.keras.callbacks.TensorBoard(log_dir=tensorboard_path,
histogram_freq=1)
callbacks.append(tb_callback)
model.fit(x=X_train,
y=y_train,
epochs=epochs,
validation_data=validation_data,
callbacks=callbacks,
verbose=False)
return model
def train(model_rs: tf.keras.Model, model_kge: tf.keras.Model, train_data: List[Tuple[int, int, int]],
test_data: List[Tuple[int, int, int]], kg: List[Tuple[int, int, int]], topk_data: TopkData,
optimizer_rs=None, optimizer_kge=None, kge_interval=3, epochs=100, batch=512):
if optimizer_rs is None:
optimizer_rs = tf.keras.optimizers.Adam()
if optimizer_kge is None:
optimizer_kge = tf.keras.optimizers.Adam()
def xy(data):
user_id = tf.constant([d[0] for d in data], dtype=tf.int32)
item_id = tf.constant([d[1] for d in data], dtype=tf.int32)
head_id = tf.constant([d[1] for d in data], dtype=tf.int32)
label = tf.constant([d[2] for d in data], dtype=tf.float32)
return {'user_id': user_id, 'item_id': item_id, 'head_id': head_id}, label
def xy_kg(kg):
item_id = tf.constant([d[0] for d in kg], dtype=tf.int32)
head_id = tf.constant([d[0] for d in kg], dtype=tf.int32)
relation_id = tf.constant([d[1] for d in kg], dtype=tf.int32)
tail_id = tf.constant([d[2] for d in kg], dtype=tf.int32)
label = tf.constant([0] * len(kg), dtype=tf.float32)
return {'item_id': item_id, 'head_id': head_id, 'relation_id': relation_id, 'tail_id': tail_id}, label
train_ds = tf.data.Dataset.from_tensor_slices(xy(train_data)).shuffle(len(train_data)).batch(batch)
test_ds = tf.data.Dataset.from_tensor_slices(xy(test_data)).batch(batch)
kg_ds = tf.data.Dataset.from_tensor_slices(xy_kg(kg)).shuffle(len(kg)).batch(batch)
model_rs.compile(optimizer=optimizer_rs, loss='binary_crossentropy', metrics=['AUC', 'Precision', 'Recall'])
model_kge.compile(optimizer=optimizer_kge, loss=lambda y_true, y_pre: y_pre)
for epoch in range(epochs):
model_rs.fit(train_ds, epochs=epoch + 1, verbose=0, validation_data=test_ds,
callbacks=[RsCallback(topk_data, _get_score_fn(model_rs))], initial_epoch=epoch)
if epoch % kge_interval == 0:
model_kge.fit(kg_ds, epochs=epoch + 1, verbose=0, callbacks=[_KgeCallback()], initial_epoch=epoch)
def train_td_value_model(model: tf.keras.Model) -> None:
'''Trains a compiled model.
The model is trained using TD-learning: it tries to predict 1 plus the
maximum of its own prediction on the successor states.
'''
# TODO: we generate the labels using the model that we train. Consider
# keeping the label model fixed for a while until the trained model
# converges, and only then updating it.
examples = get_td_value_examples(model).shuffle(
buffer_size=4096).batch(BATCH_SIZE).prefetch(16)
evaluation_df = pd.DataFrame()
for epoch in range(NUM_EPOCHS):
model.fit(
x=examples,
epochs=epoch + 1, # Train for one epoch.
steps_per_epoch=NUM_STEPS_PER_EPOCH,
initial_epoch=epoch,
callbacks=[
tf.keras.callbacks.TensorBoard(log_dir='/tmp/tensorboard')
])
epoch_evaluation_df = evaluate_model(model)
print(epoch_evaluation_df)
epoch_evaluation_df['epoch'] = epoch
evaluation_df = evaluation_df.append(epoch_evaluation_df,
ignore_index=True)
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 train(self,
model: tf.keras.Model,
train_dataset: tf.data.Dataset,
steps_per_epoch: int,
val_dataset: Optional[tf.data.Dataset],
validation_steps: int,
epochs: Optional[int] = None,
batch_size: Optional[int] = None,
val_json_file: Optional[str] = None) -> tf.keras.Model:
"""Run EfficientDet training."""
config = self.config
if not epochs:
epochs = config.num_epochs
if not batch_size:
batch_size = config.batch_size
config.update(
dict(
steps_per_epoch=steps_per_epoch,
eval_samples=batch_size * validation_steps,
val_json_file=val_json_file,
batch_size=batch_size))
train.setup_model(model, config)
train.init_experimental(config)
model.fit(
train_dataset,
epochs=epochs,
steps_per_epoch=steps_per_epoch,
callbacks=train_lib.get_callbacks(config.as_dict(), val_dataset),
validation_data=val_dataset,
validation_steps=validation_steps)
return model
def _train_bert_multitask_keras_model(
train_dataset: tf.data.Dataset,
eval_dataset: tf.data.Dataset,
model: tf.keras.Model,
params: BaseParams,
mirrored_strategy: tf.distribute.MirroredStrategy = None):
# can't save whole model with model subclassing api due to tf bug
# see: https://github.com/tensorflow/tensorflow/issues/42741
# https://github.com/tensorflow/tensorflow/issues/40366
model_checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(
filepath=os.path.join(params.ckpt_dir, 'model'),
save_weights_only=True,
monitor='val_mean_acc',
mode='auto',
save_best_only=True)
tensorboard_callback = tf.keras.callbacks.TensorBoard(
log_dir=params.ckpt_dir)
if mirrored_strategy is not None:
with mirrored_strategy.scope():
model.fit(
x=train_dataset.repeat(),
validation_data=eval_dataset,
epochs=params.train_epoch,
callbacks=[model_checkpoint_callback, tensorboard_callback],
steps_per_epoch=params.train_steps_per_epoch)
else:
model.fit(x=train_dataset.repeat(),
validation_data=eval_dataset,
epochs=params.train_epoch,
callbacks=[model_checkpoint_callback, tensorboard_callback],
steps_per_epoch=params.train_steps_per_epoch)
model.summary()
def train_neural_network(
trainset: tf.data.Dataset,
validset: tf.data.Dataset,
neural_network: tf.keras.Model,
max_epochs: int,
steps_per_epoch: int,
neural_network_name: Union[str, None] = None) -> tf.keras.Model:
print('')
print('Structure of neural network:')
print('')
neural_network.summary()
print('')
callbacks = [
tf.keras.callbacks.EarlyStopping(patience=3,
monitor='val_auc',
mode='max',
restore_best_weights=True,
verbose=1)
]
if neural_network_name is not None:
callbacks.append(
tf.keras.callbacks.ModelCheckpoint(filepath=neural_network_name,
monitor='val_auc',
mode='max',
save_best_only=True,
save_weights_only=True))
neural_network.fit(trainset,
validation_data=validset,
steps_per_epoch=steps_per_epoch,
epochs=max_epochs,
verbose=1,
callbacks=callbacks)
print('')
return neural_network
def train_model(
train_generator: Iterable[Tuple[np.ndarray, np.ndarray]],
model: tf.keras.Model,
optimizer: tf.keras.optimizers.Optimizer,
loss: tf.keras.losses.Loss,
epochs: int,
val_generator: Iterable[Tuple[np.ndarray, np.ndarray]],
metrics: List[tf.keras.metrics.Metric],
callbacks: List[tf.keras.callbacks.Callback],
path_to_save_results: str,
class_weights: Optional[Dict[int, float]] = None,
loss_weights: Optional[Dict['str', float]] = None) -> tf.keras.Model:
# create directory for saving results
if not os.path.exists(path_to_save_results):
os.makedirs(path_to_save_results)
# compile model
model.compile(optimizer=optimizer,
loss=loss,
metrics=metrics,
loss_weights=loss_weights)
model.fit(train_generator,
epochs=epochs,
callbacks=callbacks,
validation_data=val_generator,
verbose=2,
class_weight=class_weights)
return model
def train_value_model(model: tf.keras.Model,
examples: tf.data.Dataset) -> None:
'''Takes a compiled model and trains it.
Trains a value model in a supervised mannel: it simply tries to predict the
value at each state.
'''
# Training examples are generated from trajectories, so consecutive
# examples are strongly correlated. This increases the variance of the
# gradient. Shuffling the examples reduces the variance and speeds up
# training significantly.
examples = examples.shuffle(
buffer_size=4096).batch(BATCH_SIZE).prefetch(16)
evaluation_df = pd.DataFrame()
for epoch in range(NUM_EPOCHS):
model.fit(
x=examples,
epochs=epoch + 1, # Train for one epoch.
steps_per_epoch=NUM_STEPS_PER_EPOCH,
initial_epoch=epoch,
callbacks=[
tf.keras.callbacks.TensorBoard(log_dir='/tmp/tensorboard')
])
epoch_evaluation_df = evaluate_model(model)
print(epoch_evaluation_df)
epoch_evaluation_df['epoch'] = epoch
evaluation_df = evaluation_df.append(epoch_evaluation_df,
ignore_index=True)
def keras_fit(
model: tf.keras.Model,
dataset: tf.data.Dataset,
num_epochs: int,
batch_size: int,
callbacks: List[tf.keras.callbacks.Callback],
) -> None:
"""Train the model using model.fit(...)."""
ds_train = dataset.batch(batch_size=batch_size, drop_remainder=False)
model.fit(ds_train, epochs=num_epochs, callbacks=callbacks, verbose=2)
def train_keras_model(model: tf.keras.Model, input_config, train_config,
task_config: TrainTaskConfig):
example_config = input_config['example_config']
label_col = input_config['label_col']
feature_specs = parse_feature_specs(example_config)
dataset_fn = get_dataset_fn(feature_specs=feature_specs,
label_col=label_col,
**train_config)
model.fit(dataset_fn())
def train(epochs: int, data_handler: DataHandler, model: tf.keras.Model, models_dir: str, model_name: str = "unnamed",
batch_size: int = 64, loss_weights: dict = None):
for epoch in range(epochs):
print(f"Starting Epoch {epoch}")
train_data = data_handler.train_iterator(batch_size=batch_size)
model.fit(train_data, class_weight=loss_weights)
print(f"Finished training on Epoch {epoch}")
test_data = data_handler.test_iterator(batch_size=batch_size)
model.evaluate(test_data)
print(f"Finished evaluation on Epoch {epoch}")
model.save(os.path.join(models_dir, model_name, "checkpoints", str(epoch)))
model.save(os.path.join(model_name, model_name, "final_model"))
def train(model: tf.keras.Model, training_data, testing_data, callbacks=None):
# Check for checkpoint
if callbacks:
history = model.fit(training_data, epochs=EPOCHS,
validation_data=testing_data,
validation_steps=VALIDATION_STEPS,
callbacks=callbacks)
else:
history = model.fit(training_data, epochs=EPOCHS,
validation_data=testing_data,
validation_steps=VALIDATION_STEPS)
return history
def train_and_evaluate_model(model: tf.keras.Model) -> None:
"""Train and test the transfer learning model
Parameters
----------
model : tf.keras.Model
The transfer learning model
"""
optimizer = tf.keras.optimizers.Adam(learning_rate=1e-5)
loss = tf.keras.losses.BinaryCrossentropy(from_logits=True)
metric = tf.keras.metrics.BinaryAccuracy()
model.compile(loss=loss, optimizer=optimizer, metrics=[metric])
train_dataset, validation_dataset, test_dataset = get_dataset()
model.fit(train_dataset, epochs=20, validation_data=validation_dataset)
model.evaluate(x=test_dataset)
def train_model(model: tf.keras.Model,
tr_dataset: tf.data.Dataset,
val_dataset: tf.data.Dataset,
checkpoint_dir: str,
epochs: int = 100,
patience: int = 10,
save_one: bool = False):
def loss(labels, logits):
return tf.keras.losses.sparse_categorical_crossentropy(
labels, logits, from_logits=True)
optimizer = tf.keras.optimizers.Adam()
model.compile(optimizer=optimizer, loss=loss)
early_stop = tf.keras.callbacks.EarlyStopping(monitor='val_loss',
patience=patience)
logger.info("Begin training... (this will take a while)")
checkpoint_prefix = os.path.join(checkpoint_dir,
"ckpt" if save_one else "ckpt_{epoch}")
checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(
filepath=checkpoint_prefix,
save_best_only=save_one,
save_weights_only=True)
history = model.fit(tr_dataset,
epochs=epochs,
callbacks=[checkpoint_callback, early_stop],
validation_data=val_dataset)
logger.info(
"Training stopped, no improvement after {} epochs".format(patience))
def fit_loop(
self,
model: tf.keras.Model,
Xy: Sequence[Tuple[np.ndarray, np.ndarray]],
validation_data: Optional[Tuple[np.ndarray, np.ndarray]] = None,
callbacks: Iterable[Callable[[EpochResult], None]] = (),
) -> None:
"""
Args:
model: keras model to train
Xy: sequence of data batches to be passed to `model.fit`
validation_data: passed as `validation_data` argument to `model.fit`
callbacks: if given, these will be called at the end of each epoch
"""
for i_epoch in range(self.epochs):
Xy = shuffle(Xy)
if self.workers > 1:
Xy = ThreadedSequencePreLoader(
Xy,
num_workers=self.workers,
max_queue_size=self.max_queue_size)
history = []
for i_batch, (X, y) in enumerate(Xy):
logger.info(f"Fitting on batch {i_batch + 1} of {len(Xy)}, "
f"of epoch {i_epoch}...")
history.append(
model.fit(
X,
y,
validation_data=validation_data,
batch_size=self.batch_size,
))
for callback in callbacks:
callback(EpochResult(epoch=i_epoch, history=tuple(history)))
def train(
model: tf.keras.Model,
img_shape: Tuple[Optional[int], Optional[int], Optional[int]],
checkpoint_path: Path,
epochs: int = 20,
batch_size: int = 8,
save_freq: int = 5,
):
(ds_train, ds_valid) = tfds.load( # NOQA (false positve)
"mura",
split=["train", "valid"],
shuffle_files=True,
as_supervised=True,
)
ds_train = _prepare_ds(ds_train,
img_shape=img_shape,
batch_size=batch_size)
ds_valid = _prepare_ds(ds_valid,
img_shape=img_shape,
batch_size=batch_size)
callbacks = _train_calbacks(checkpoint_path, save_freq)
return model.fit(
ds_train,
epochs=epochs,
validation_data=ds_valid,
callbacks=callbacks,
)
def train(model: tf.keras.Model) -> tf.keras.Model:
"""Trains the classification model on the chess piece image dataset.
"""
model.compile(
optimizer="adam",
loss=losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=["accuracy"],
)
train_dataset = _get_dataset("training")
validation_dataset = _get_dataset("validation")
model.fit(train_dataset,
validation_data=validation_dataset,
epochs=TRAINING_EPOCHS)
return model
def train_model(model: tf.keras.Model,
train_data,
validation_data,
optimizer,
loss='categorical_crossentropy',
epochs=3,
verbose=1,
batch_size=None,
callbacks=None):
# init
K.clear_session()
tf.random.set_seed(51)
np.random.seed(51)
# optimizer
opt = tf.keras.optimizers.Adam() if optimizer is None else optimizer
# compile
model.compile(opt, loss=loss, metrics=["acc"])
# fit
history = model.fit(train_data,
validation_data=validation_data,
epochs=epochs,
verbose=verbose,
callbacks=callbacks,
batch_size=batch_size)
return history
def train(model: tf.keras.Model, data_csv: str, data_root: str, hparams: dict,
hp_space: dict, run_name: str):
train_split = hparams[hp_space['train_split']]
num_imgs = hparams[hp_space['num_imgs']]
num_epochs = hparams[hp_space['num_epochs']]
batch_size = hparams[hp_space['batch_size']]
# Load csv data
examples_df = pd.read_csv(data_csv, header=0, skipinitialspace=True)
label_dict = {'Fish': 0, 'Flower': 1, 'Gravel': 2, 'Sugar': 3}
filenames = np.array(
[data_root + fname for fname in examples_df['Image'].tolist()],
dtype=object)
labels = np.array(
[label_dict[fname] for fname in examples_df['Label'].tolist()],
dtype=int)
if num_imgs == -1: num_imgs = len(filenames)
train_idxs, val_idxs = split_train_val(examples_df, train_split, num_imgs)
train_filenames = filenames[train_idxs]
train_labels = labels[train_idxs]
val_filenames = filenames[val_idxs]
val_labels = labels[val_idxs]
train_dataset, num_train = create_dataset(files=train_filenames,
labels=train_labels,
batch_size=batch_size,
training=True)
val_dataset, num_val = create_dataset(files=val_filenames,
labels=val_labels,
batch_size=batch_size,
training=False)
print("Num train: {}, num val: {}".format(num_train, num_val))
log_dir = './logs/{}'.format(run_name)
train_history = model.fit(
train_dataset,
epochs=num_epochs,
steps_per_epoch=num_train // batch_size,
validation_steps=num_val // batch_size,
validation_data=val_dataset,
callbacks=[
tf.keras.callbacks.ReduceLROnPlateau(monitor='val_loss',
factor=0.5,
patience=5),
tf.keras.callbacks.TensorBoard(log_dir=log_dir),
tf.keras.callbacks.ModelCheckpoint(
filepath="experiments/{}/".format(run_name) +
"best_model.hdf5",
save_best_only=True,
save_weights_only=False),
hp.KerasCallback(log_dir, hparams)
])
def fit(
model: tf.keras.Model,
train_data,
validation_data=None,
epochs: int = 1,
initial_epoch: int = 0,
validation_freq: int = 1,
callbacks: Iterable[tf.keras.callbacks.Callback] = (),
steps_per_epoch: Optional[int] = None,
verbose: bool = True,
jit_compile: bool = False,
):
"""
Call `fit_single` or `Model.fit` based on `train_data`.
Delegates to either `graph_tf.train.fit_single` or `tf.keras.Model.fit`.
Args:
model: keras model to train.
train_data: (inputs, labels, sample_weight) or dataset with a
single element for training.
validation_data: (inputs, labels, sample_weight) or dataset with a
single element for validation.
epochs: int, maximum number of steps/epochs to train for.
initial_epoch: int, starting epoch.
validation_freq: int, number of training steps/epochs per validation.
callbacks: Iterable of `tf.keras.callbacks.Callbacks`.
steps_per_epoch: Number of steps per epoch. Must be 1 if specified and
train_data is a not a `tf.data.Dataset`.
verbose: flag resulting in verbose outputs.
jit_compile: used in fit_single. Ignored if more than one example.
Returns:
history: `tf.keras.callbacks.History` object.
"""
if not isinstance(train_data, tf.data.Dataset) or len(train_data) == 1:
assert steps_per_epoch is None or steps_per_epoch == 1
return fit_single(
model=model,
train_data=train_data,
validation_data=validation_data,
epochs=epochs,
initial_epoch=initial_epoch,
validation_freq=validation_freq,
callbacks=callbacks,
verbose=verbose,
jit_compile=jit_compile,
)
return model.fit(
train_data,
validation_data=validation_data,
epochs=epochs,
initial_epoch=initial_epoch,
validation_freq=validation_freq,
callbacks=callbacks,
verbose=verbose,
steps_per_epoch=steps_per_epoch,
)
def train_model(model: tf.keras.Model,
train_input: np.ndarray,
train_output: np.ndarray,
batch_size: int = 8,
epochs: int = 10,
validation_split: float = 0.1,
model_path: str = "model.h5",
**kwargs):
""" trains and saves the model
"""
model.fit(train_input,
train_output,
batch_size=batch_size,
epochs=epochs,
validation_split=validation_split)
print(model.summary())
model.save(model_path)
def test_can_run_the_lr_finder(model: tf.keras.Model,
dataset: tf.data.Dataset):
min_lr = 1e-6
max_lr = 1e-1
model.compile(optimizer=tf.keras.optimizers.SGD(),
loss=tf.keras.losses.MeanSquaredError())
lrfinder = LRFinder(min_lr, max_lr, num_steps=NUM_STEPS, figname=FIGNAME)
model.fit(dataset, epochs=1, callbacks=[lrfinder])
assert len(lrfinder.losses) == NUM_STEPS
assert len(lrfinder.lrs) == NUM_STEPS
assert lrfinder.lrs[0] == min_lr
assert lrfinder.lrs[-1] == max_lr
# by default should have saved a figure with the results
assert os.path.exists(lrfinder.figname)
def train_model(model: tf.keras.Model, x: np.ndarray, y: np.ndarray,
epochs: int, batch_size: int) -> tf.keras.callbacks.History:
"""Trains the given binary classification model."""
history = model.fit(x=x,
y=y,
batch_size=batch_size,
epochs=epochs,
validation_split=0.2,
verbose=0)
return history
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 compile_and_fit(model: tf.keras.Model, window: WindowGenerator,
model_name: str,
patience=cfg.EARLY_STOPPING['patience'],
):
"""
Train model
@param model_name:
@param model:
@param window:
@param patience:
@return:
"""
checkpoint_dir = os.path.join(cfg.CHECKPOINT_PATH, model_name)
checkpoint_path = os.path.join(checkpoint_dir, '{epoch:04d}.ckpt')
make_dir(checkpoint_dir)
callbacks = []
if not is_dir_empty(checkpoint_dir):
load_weight(model, checkpoint_dir)
cp_callback = tf.keras.callbacks.ModelCheckpoint(
checkpoint_path,
save_weights_only=True,
verbose=1,
)
callbacks.append(cp_callback)
if cfg.EARLY_STOPPING['enabled'] is True:
early_stopping = tf.keras.callbacks.EarlyStopping(
monitor='val_loss',
patience=patience,
mode='min')
callbacks.append(early_stopping)
model.compile(
optimizer=tf.keras.optimizers.Adam(),
loss=tf.losses.MeanSquaredError(),
metrics=[tf.metrics.MeanAbsoluteError()]
)
history = model.fit(
window.train,
epochs=cfg.MAX_EPOCH,
validation_data=window.val,
callbacks=callbacks,
verbose=2,
)
return history
def train(args: Settings, model: tf.keras.Model, dataset: tf.data.Dataset,
path):
callbacks = [
# tf.keras.callbacks.EarlyStopping(patience=2),
tf.keras.callbacks.ModelCheckpoint(filepath=str(path['ckpt'] /
'{epoch:03d}.ckpt')),
tf.keras.callbacks.TensorBoard(
update_freq=args.update_freq, # every 128 batches
log_dir=path['log'],
profile_batch='2,66'),
ShowImageCallback(args,
log_dir=path['log'] / 'flow',
log_period=args.update_freq)
]
optimizer = tf.keras.optimizers.Adam(learning_rate=args.learning_rate)
losses = [AutoResizeMseLoss() for _ in model.outputs]
model.compile(
optimizer=optimizer,
loss=losses,
# metrics=[tf.keras.metrics.MeanSquaredError()]
)
if args.load_ckpt:
weight_path = Path(args.load_ckpt) / 'variables/variables'
model.load_weights(weight_path, by_name=False)
try:
model.fit(dataset, epochs=args.num_epoch, callbacks=callbacks)
except KeyboardInterrupt as e:
pass
finally:
for ext, fmt in zip(['pb', 'h5'], ['tf', 'h5']):
out_file = str(path['run'] / 'model.{}'.format(ext))
logging.info(
'saving weights to {} prior to termination ...'.format(
out_file))
model.save_weights(out_file, save_format=fmt)
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 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
