Load HDF5 checkpoint models with custom metrics - tensorflow

Question

I've created a Keras Regressor to run a RandomizedSearch CV using a ModelCheckpoint callback, but the training overran the Colab runtime 12H limit and stopped halfway through. The models are saved in hdf5 format.
I used tensorflow_addons to add the RSquare class to monitor the R2 for train and validation sets. However, when I used keras.models.load_model, I get the following error:
As you can see from the traceback, I have passed the custom_objects parameter, but still it is not recognised.
How can I solve this?
You can see the full code example below:
import os
import tensorflow as tf
import tensorflow_addons as tfa
from tensorflow.keras import Sequential
from tensorflow.keras.layers import Dense, Dropout
from tensorflow.keras.layers import Input, InputLayer
from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint
from keras.wrappers.scikit_learn import KerasRegressor
from sklearn.model_selection import RandomizedSearchCV
def build_model(n_hidden = 2, n_neurons = 64, input_shape = (X_train.shape[1],), dropout = 0):
model = Sequential()
model.add(InputLayer(input_shape = input_shape))
for i in range(n_hidden):
model.add(Dense(n_neurons, activation = 'relu'))
model.add(Dropout(dropout))
model.add(Dense(1))
model.compile(loss = 'mean_squared_error', optimizer = 'adam', metrics = [tfa.metrics.RSquare(y_shape=(1,))])
return model
keras_reg = KerasRegressor(build_model)
checkdir = os.path.join(r'/content/drive/MyDrive/COP328/Case A', 'checkpoints', datetime.datetime.now().strftime('%d-%m-%Y_%H-%M-%S'), 'imputed_log1p-{epoch:02d}-{val_r_square:.3f}.hdf5')
callbacks = [ModelCheckpoint(checkdir, save_freq='epoch', save_best_only = True, monitor = 'val_r_square', mode = 'max'),
EarlyStopping(patience = 10)]
# Here is where training got interrupted because of Colab runtime being dropped:
param_dist = {
'n_hidden' : [1,2],
'n_neurons': [8,16,32,64,128],
'dropout': [0,0.2,0.4]
}
rnd_search_cv = RandomizedSearchCV(keras_reg, param_dist, n_iter= 15, cv = 5)
rnd_search_cv.fit(X_train, y_train, epochs = 200, batch_size = 64,
validation_data = (X_valid,y_valid),
callbacks = callbacks)
# Here is where I am trying to reload one of the most promising models based on R2, and getting the error:
from keras.models import load_model
import tensorflow_addons as tfa
rnd_model = load_model(r'/content/drive/MyDrive/COP328/Case A/checkpoints/26-06-2021_17-32-29/imputed_log1p-56-1.000.hdf5', custom_objects = {'r_square': tfa.metrics.RSquare, 'val_r_square': tfa.metrics.RSquare})

Answer 1

This solution doesn't take into account the addons package but one possibility is to create the coefficient of determination (R^2) as a different metric without that package, and then defining it as your loss.
def coeff_determination(y_true, y_pred):
from keras import backend as K
SS_res = K.sum(K.square( y_true-y_pred ))
SS_tot = K.sum(K.square( y_true - K.mean(y_true) ) )
return ( 1 - SS_res/(SS_tot + K.epsilon()) )
Recall that minimizing MSE will maximize R^2.