Issue when using Keras callback with tf.keras

[fmannhardt]

I get the following error when using the Keras callback with tf.keras:

algo-1-bvae4_1 | Traceback (most recent call last):
algo-1-bvae4_1 | File "model.py", line 211, in
algo-1-bvae4_1 | verbose = 2)
algo-1-bvae4_1 | File "/usr/local/lib/python3.6/site-packages/tensorflow/python/keras/engine/training.py", line 780, in fit
algo-1-bvae4_1 | steps_name='steps_per_epoch')
algo-1-bvae4_1 | File "/usr/local/lib/python3.6/site-packages/tensorflow/python/keras/engine/training_arrays.py", line 267, in model_iteration
algo-1-bvae4_1 | callbacks._call_batch_hook(mode, 'begin', step, batch_logs)
algo-1-bvae4_1 | File "/usr/local/lib/python3.6/site-packages/tensorflow/python/keras/callbacks.py", line 247, in _call_batch_hook
algo-1-bvae4_1 | batch_hook = getattr(callback, hook_name)
algo-1-bvae4_1 | AttributeError: 'LosswiseKerasCallback' object has no attribute 'on_train_batch_begin'

I assume those are not compatible? Or does one simply need to add the on_train_batch_begin method?

[fmannhardt]

It seems that extending it trivially like this solves the issue:

class LosswiseKerasTFCallback(LosswiseKerasCallback):
    
    def __init__(self, tag=None, params={}, track_git=False, display_interval=None, max_iter=None):
        super(LosswiseKerasTFCallback, self).__init__(tag, params, track_git, display_interval, max_iter)
    
    def on_train_batch_begin(self, batch, logs={}):
        return()
    
    def on_train_batch_end(self, batch, logs={}):
        self.on_batch_end(batch, logs)

Maybe this helps someone else tripping over this issue.