Customized TF2 Model Save

Question

I write customize model with TF2

class NN(tf.keras.Model):

def __init__(self,
             output_dim: int, 
             controller_dime:int=128,
             interface_dim: int=35,
             netsize: int=100, 
             degree: int=20, 
             k:float=2,
             name:str='dnc_rn')->None:

---

Its full of random parametrs that is not trainable! so I need to save model completely, and I cant use save_weights because train of each model dependes on its self random parameters...

---

the file of trainer is like :

import numpy as np
import tensorflow as tf

def trainer(model: tf.keras.Model,
        loss_fn: tf.keras.losses,
        X_train: np.ndarray,
        y_train: np.ndarray = None,
        optimizer: tf.keras.optimizers = tf.keras.optimizers.Adam(learning_rate=1e-3),
        loss_fn_kwargs: dict = None,
        epochs: int = 1000000,
        batch_size: int = 1,
        buffer_size: int = 2048,
        shuffle: bool = False,
        verbose: bool = True,
        show_model_interface_vector: bool = False
        ) -> None:

"""
Train TensorFlow model.

Parameters
----------
model
    Model to train.
loss_fn
    Loss function used for training.
X_train
    Training batch.
y_train
    Training labels.
optimizer
    Optimizer used for training.
loss_fn_kwargs
    Kwargs for loss function.
epochs
    Number of training epochs.
batch_size
    Batch size used for training.
buffer_size
    Maximum number of elements that will be buffered when prefetching.
shuffle
    Whether to shuffle training data.
verbose
    Whether to print training progress.
"""

model.show_interface_vector=show_model_interface_vector

# Create dataset
if y_train is None:  # Unsupervised model
    train_data = X_train
else:
    train_data = (X_train, y_train)
train_data = tf.data.Dataset.from_tensor_slices(train_data)
if shuffle:
    train_data = train_data.shuffle(buffer_size=buffer_size).batch(batch_size)

# Iterate over epochs
history=[]
for epoch in range(epochs):
    if verbose:
        pbar = tf.keras.utils.Progbar(target=epochs, width=40, verbose=1, interval=0.05)

    # Iterate over the batches of the dataset
    for step, train_batch in enumerate(train_data):

        if y_train is None:
            X_train_batch = train_batch
        else:
            X_train_batch, y_train_batch = train_batch

        with tf.GradientTape() as tape:
            preds = model(X_train_batch)

            if y_train is None:
                ground_truth = X_train_batch
            else:
                ground_truth = y_train_batch

            # Compute loss
            if tf.is_tensor(preds):
                args = [ground_truth, preds]
            else:
                args = [ground_truth] + list(preds)

            if loss_fn_kwargs:
                loss = loss_fn(*args, **loss_fn_kwargs)
            else:
                loss = loss_fn(*args)

            if model.losses:  # Additional model losses
                loss += sum(model.losses)

        grads = tape.gradient(loss, model.trainable_weights)
        optimizer.apply_gradients(zip(grads, model.trainable_weights))

    if verbose:
            loss_val = loss.numpy().mean()
            pbar_values = [('loss', loss_val)]
            pbar.update(epoch+1, values=pbar_values)

    history.append(loss.numpy().mean())

model.show_interface_vector= not show_model_interface_vector
return history

---

after training very I tried to save model but when I call TF2 .save :

model.save('a.h5')

I have an error:

NotImplementedError: Saving the model to HDF5 format requires the model to be a Functional model or a Sequential model. It does not work for subclassed models, because such models are defined via the body of a Python method, which isn't safely serializable. Consider saving to the Tensorflow SavedModel format (by setting save_format="tf") or using `save_weights`.

I change it to .tf format but again:

ValueError: Model <model2.NN object at 0x11448b390> cannot be saved because the input shapes have not been set. Usually, input shapes are automatically determined from calling .fit() or .predict(). To manually set the shapes, call model._set_inputs(inputs).

but its already trained, and if i _set_inputs

ValueError: Cannot infer num from shape (None, 12, 4)

---

I don't know what should I do? I am sciences and amateur with TF2 Help me it's important for my project...

Answer 1

The first error says that you can't use model.save for subclassed models. Subclassed models are of the type: you define a class which inherits from tf.keras.models.Model. As the error suggests, try using functional or sequential api to construct your model so you can save in h5 format.

Answer 2

I found the answer and just using:

import pickle
import dill

dill.dump(model, file = open("model.pickle", "wb"))
model_reloaded = dill.load(open("model.pickle", "rb"))

from this topic: Saving an Object (Data persistence)

Answer 3

If you use your self-defined complex model which means you create an optimizer, compute gradients and apply gradients to some complex parts, tf.keras.Model.save is not suitable especially when the shape of inputs is undefined (just I think).

So the tf.train.Checkpoint API is suitable for such a situation. Click the link for the tutorial. The tf.train.Checkpoint can save both the model and the optimizer at the same time. The way of using tf.train.Checkpoint is similar to the way of using tf.compat.v1.train.Saver. Thus it may be the alternative for migrating code from tensorflow1 to tensorflow2.