Following this guide https://keras.io/guides/customizing_what_happens_in_fit/ I have created a custom version of the train_step that will be called when calling model.fit()
#tf.function
def train_step(self, x, y):
''' Adapted from https://keras.io/guides/customizing_what_happens_in_fit/ '''
with tf.GradientTape() as tape:
logits = model(x, training=True)
trainable_vars = model.trainable_variables
loss = custom_loss(y, logits)
gradients = tape.gradient(loss, trainable_vars)
self.optimizer.apply_gradients(zip(gradients, trainable_vars))
self.compiled_metrics.update_state(y, logits)
return {m.name: m.result() for m in self.metrics}
when I then fit the model as
model.compile(optimizer=optimizer, loss=custom_loss)
model.fit(x,y)
the model is trained, however when I try to print something inside the train_step() function I see no output, making me wonder if .fit() actually calls my customized function? Or is what I'm trying to print being printed somewhere in the backend that I cant see in the console?
you should better use tf.print("...") inside train_step during training. I suggest to keep the original tensorflow fit method by replacing it with callback one, then you have multiple options allow you to run your code during training or at the end of epoch or validation.
Example of creating customized callback which print date after the beginning and the end of each epoch:
import tensorflow as tf
from datetime import datetime
class date_callback(tf.keras.callbacks.Callback):
def __init__(self):
super(date_callback, self).__init__()
#other inits if necessary
def on_epoch_begin(self, epoch, logs = None):
now = datetime.now()
tf.print("epoch begin at", now)
def on_epoch_end(self, epoch, logs = None):
now = datetime.now()
tf.print("epoch end at", now)
...
...
#running the model training
dc = date_callback()
model.fit(..., callbacks = [dc])
Related
I am using the following code to make a custom train step: https://keras.io/guides/customizing_what_happens_in_fit/
I'd like to use this with tf.distribute.MultiWorkerMirroredStrategy. How should the loss be calculated if using a custom loss? As an example, let's say my custom loss is tf.keras.losses.SparseCategoricalCrossentropy. Should the reduction be set to tf.keras.losses.Reduction.NONE (like below)? Or do I need to set the reduction to tf.keras.losses.Reduction.SUM and divide by the global batch size?
My initial thought is to set the reduction to tf.keras.losses.Reduction.NONE. Then when I call model.fit, keras will automatically handle gradient aggregation across the replicas.
class CustomModel(tf.keras.Model):
def __init__(self, model):
super(CustomModel, self).__init__()
self.model = model
self.loss_tracker= tf.keras.metrics.Mean(name='loss')
def train_step(self, data):
x, y = data
with tf.GradientTape() as tape:
# Forward pass
y_pred = self.model(x, training=True)
# Compute our own loss. Shape (batch_size,)
loss = tf.keras.losses.SparseCategoricalCrossentropy(
from_logits = True, reduction=tf.keras.losses.Reduction.NONE)(y, y_pred)
# Compute gradients
trainable_vars = self.model.trainable_variables
gradients = tape.gradient(loss, trainable_vars)
# Update weights
self.optimizer.apply_gradients(zip(gradients, trainable_vars))
# Update metrics
self.loss_tracker.update_state(loss)
self.compiled_metrics.update_state(y, y_pred)
metrics = {m.name: m.result() for m in self.metrics}
return metrics
According to the tensorflow documentation, it should be possible, to configure the loss inside the custom train_step function and only configure the optimizer in compile().
My class looks like this:
import tensorflow as tf
from tensorflow import keras
loss_tracker = tf.keras.metrics.Mean(name="loss")
loss_fn = keras.losses.SparseCategoricalCrossentropy(
reduction=tf.keras.losses.Reduction.NONE
)
class MaskedLanguageModel(tf.keras.Model):
def train_step(self, inputs):
if len(inputs) == 3:
features, labels, sample_weight = inputs
else:
features, labels = inputs
sample_weight = None
with tf.GradientTape() as tape:
predictions = self(features, training=True)
loss = loss_fn(labels, predictions, sample_weight=sample_weight)
# Compute gradients
trainable_vars = self.trainable_variables
gradients = tape.gradient(loss, trainable_vars)
# Update weights
self.optimizer.apply_gradients(zip(gradients, trainable_vars))
# Compute our own metrics
loss_tracker.update_state(loss, sample_weight=sample_weight)
# Return a dict mapping metric names to current value
return {"loss": loss_tracker.result()}
#property
def metrics(self):
# We list our `Metric` objects here so that `reset_states()` can be
# called automatically at the start of each epoch
# or at the start of `evaluate()`.
# If you don't implement this property, you have to call
# `reset_states()` yourself at the time of your choosing.
return [loss_tracker]
And the model is compiled with the following snippet.
mlm_model = MaskedLanguageModel(inputs, mlm_output, name="masked_bert_model")
optimizer = keras.optimizers.Adam(learning_rate=lr)
mlm_model.compile(optimizer=optimizer)
But I get the following error:
'The model cannot be compiled because it has no loss to optimize.'
(My tensorflow version is 2.1.0)
Thank you in advance for any hint. :)
I am using Pytorch Lightning to train my models (on GPU devices, using DDP) and TensorBoard is the default logger used by Lightning.
My code is setup to log the training and validation loss on each training and validation step respectively.
class MyLightningModel(pl.LightningModule):
def training_step(self, batch):
x, labels = batch
out = self(x)
loss = F.mse_loss(out, labels)
self.log("train_loss", loss)
return loss
def validation_step(self, batch):
x, labels = batch
out = self(x)
loss = F.mse_loss(out, labels)
self.log("val_loss", loss)
return loss
TensorBoard correctly plots both the train_loss and val_loss charts in the SCALERS tab. However, in the HPARAMS tab, on the left side bar, only hp_metric is visible under Metrics.
However, in the HPARAMS tab, on the left side bar, only hp_metric is visible under Metrics.
How can we add train_loss and val_loss to the Metrics section? This way, we will be able to use val_loss in the PARALLEL COORDINATES VIEW instead of hp_metric.
Image showing hp_metric and no val_loss:
Using Pytorch 1.8.1, Pytorch Lightning 1.2.6, TensorBoard 2.4.1
You can use self.logger.log_hyperparams method to log hyperparameters and metrics in tensorboard. (see pytorch lightning tensorboard docs)
The values you added by self.log will be displayed in hparam plugin if and only if you have passed same key names in metric. (see pytorch tensorboard docs)
Sample code (full code):
class BasicModule(LightningModule):
def __init__(self, lr=0.01):
super().__init__()
self.model = models.resnet18(pretrained=False)
self.criterion = nn.CrossEntropyLoss()
self.lr = lr
self.save_hyperparameters()
metric = MetricCollection({'top#1': Accuracy(top_k=1), 'top#5': Accuracy(top_k=5)})
self.train_metric = metric.clone(prefix='train/')
self.valid_metric = metric.clone(prefix='valid/')
def on_train_start(self) -> None:
# log hyperparams
self.logger.log_hyperparams(self.hparams, {'train/top#1': 0, 'train/top#5': 0, 'valid/top#1': 0, 'valid/top#5': 0})
return super().on_train_start()
def training_step(self, batch, batch_idx, optimizer_idx=None):
return self.shared_step(*batch, self.train_metric)
def validation_step(self, batch, batch_idx):
return self.shared_step(*batch, self.valid_metric)
def shared_step(self, x, y, metric):
y_hat = self.model(x)
loss = self.criterion(y_hat, y)
self.log_dict(metric(y_hat, y), prog_bar=True)
return loss
if __name__ == '__main__':
# default_hp_metric=False
logger = loggers.TensorBoardLogger('', 'lightning_logs', default_hp_metric=False)
trainer = Trainer(max_epochs=2, gpus='0,', logger=logger, precision=16)
This question is a possible duplicate of another question on stack overflow.
However, the answer doesn't indicate how to modify the optimizer learning rate using a scheduler (that could be implemented in simple python).
I'm training a tensorflow model from scratch as explained here. Hence, an optimizer is defined as: optimizer = keras.optimizers.SGD(learning_rate=1e-3), therefore, the learning rate is defined at the beginning. However, I'd like to have a learning rate scheduler such as tf.keras.optimizers.schedules.ExponentialDecay. So how can I change the optimizer learning rate from within the training loop?
Please note that I am not using the model.fit in this case.
Any help is much appreciated.
Try this code:
import tensorflow as tf
class CustomSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):
def __init__(self, drop=0.5):
super(CustomSchedule, self).__init__()
self.drop = drop
def __call__(self, step):
return tf.math.pow(self.drop, step)
def train_step(images, labels):
with tf.GradientTape() as tape:
predictions = model(images)
loss = tf.losses.mse(labels, predictions)
gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
learning_rate = CustomSchedule(0.5)
optimizer = tf.keras.optimizers.SGD(learning_rate=learning_rate)
model = tf.keras.Sequential([
tf.keras.layers.Dense(10)
])
input = tf.random.uniform([10, 10])
labels = tf.random.uniform([10, 10], 0, 10, dtype=tf.int32)
train_step(input, labels)
I would like to keep track of the gradients over tensorboard.
However, since session run statements are not a thing anymore and the write_grads argument of tf.keras.callbacks.TensorBoard is depricated, I would like to know how to keep track of gradients during training with Keras or tensorflow 2.0.
My current approach is to create a new callback class for this purpose, but without success. Maybe someone else knows how to accomplish this kind of advanced stuff.
The code created for testing is shown below, but runs into errors independently of printing a gradient value to console or tensorboard.
import tensorflow as tf
from tensorflow.python.keras import backend as K
mnist = tf.keras.datasets.mnist
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0
model = tf.keras.models.Sequential([
tf.keras.layers.Flatten(input_shape=(28, 28)),
tf.keras.layers.Dense(128, activation='relu', name='dense128'),
tf.keras.layers.Dropout(0.2),
tf.keras.layers.Dense(10, activation='softmax', name='dense10')
])
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
class GradientCallback(tf.keras.callbacks.Callback):
console = True
def on_epoch_end(self, epoch, logs=None):
weights = [w for w in self.model.trainable_weights if 'dense' in w.name and 'bias' in w.name]
loss = self.model.total_loss
optimizer = self.model.optimizer
gradients = optimizer.get_gradients(loss, weights)
for t in gradients:
if self.console:
print('Tensor: {}'.format(t.name))
print('{}\n'.format(K.get_value(t)[:10]))
else:
tf.summary.histogram(t.name, data=t)
file_writer = tf.summary.create_file_writer("./metrics")
file_writer.set_as_default()
# write_grads has been removed
tensorboard_cb = tf.keras.callbacks.TensorBoard(histogram_freq=1, write_grads=True)
gradient_cb = GradientCallback()
model.fit(x_train, y_train, epochs=5, callbacks=[gradient_cb, tensorboard_cb])
Priniting bias gradients to console (console parameter = True)
leads to: AttributeError: 'Tensor' object has no attribute 'numpy'
Writing to tensorboard (console parameter = False) creates:
TypeError: Using a tf.Tensor as a Python bool is not allowed. Use if t is not None: instead of if t: to test if a tensor is defined, and use TensorFlow ops such as tf.cond to execute subgraphs conditioned on the
value of a tensor.
To compute the gradients of the loss against the weights, use
with tf.GradientTape() as tape:
loss = model(model.trainable_weights)
tape.gradient(loss, model.trainable_weights)
This is (arguably poorly) documented on GradientTape.
We do not need to tape.watch the variable because trainable parameters are watched by default.
As a function, it can be written as
def gradient(model, x):
x_tensor = tf.convert_to_tensor(x, dtype=tf.float32)
with tf.GradientTape() as t:
t.watch(x_tensor)
loss = model(x_tensor)
return t.gradient(loss, x_tensor).numpy()
Also have a look here: https://github.com/tensorflow/tensorflow/issues/31542#issuecomment-630495970
richardwth wrote a child class of Tensorboard.
I adapted it as follows:
class ExtendedTensorBoard(tf.keras.callbacks.TensorBoard):
def _log_gradients(self, epoch):
writer = self._writers['train']
with writer.as_default(), tf.GradientTape() as g:
# here we use test data to calculate the gradients
features, y_true = list(val_dataset.batch(100).take(1))[0]
y_pred = self.model(features) # forward-propagation
loss = self.model.compiled_loss(y_true=y_true, y_pred=y_pred) # calculate loss
gradients = g.gradient(loss, self.model.trainable_weights) # back-propagation
# In eager mode, grads does not have name, so we get names from model.trainable_weights
for weights, grads in zip(self.model.trainable_weights, gradients):
tf.summary.histogram(
weights.name.replace(':', '_') + '_grads', data=grads, step=epoch)
writer.flush()
def on_epoch_end(self, epoch, logs=None):
# This function overwrites the on_epoch_end in tf.keras.callbacks.TensorBoard
# but we do need to run the original on_epoch_end, so here we use the super function.
super(ExtendedTensorBoard, self).on_epoch_end(epoch, logs=logs)
if self.histogram_freq and epoch % self.histogram_freq == 0:
self._log_gradients(epoch)