Can't save a keras model after training my model

Question

I'm following a code example from keras.io about A Vision Transformer without Attention, here. I want to save it in the tensorflow/keras after completing training but it generates an error. The last part of my model code is given below

# Get the total number of steps for training.
total_steps = int((len(x_train) / config.batch_size) * config.epochs)

# Calculate the number of steps for warmup.
warmup_epoch_percentage = 0.15
warmup_steps = int(total_steps * warmup_epoch_percentage)

# Initialize the warmupcosine schedule.
scheduled_lrs = WarmUpCosine(
lr_start=1e-5, lr_max=1e-3, warmup_steps=warmup_steps, total_steps=total_steps,
)

# Get the optimizer.
optimizer = tfa.optimizers.AdamW(
learning_rate=scheduled_lrs, weight_decay=config.weight_decay
)

# Compile and pretrain the model.
model.compile(
optimizer=optimizer,
loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=[
    keras.metrics.SparseCategoricalAccuracy(name="accuracy"),
    keras.metrics.SparseTopKCategoricalAccuracy(5, name="top-5-accuracy"),
],
)

# Train the model
history = model.fit(
train_ds,
epochs=config.epochs,
validation_data=val_ds,
callbacks=[
    keras.callbacks.EarlyStopping(monitor="val_accuracy", patience=5, mode="auto",)
],
)

# Evaluate the model with the test dataset.
print("TESTING")
loss, acc_top1, acc_top5 = model.evaluate(test_ds)
print(f"Loss: {loss:0.2f}")
print(f"Top 1 test accuracy: {acc_top1*100:0.2f}%")
print(f"Top 5 test accuracy: {acc_top5*100:0.2f}%")

I tried below two methods to save my model

model.save('/content/drive/MyDrive/VIT-SHIFT') 

and

history.save('/content/drive/MyDrive/VIT-SHIFT')

but it says model and history are not defined. Full code is available in this colab notebook

Answer 1

As mentioned in your other question about the same code, HERE, you may need to implement the call method first. Now, saving and reloading the model should be straightforward but I've encountered an issue regarding a layers.RandomCrop augmentation layer that is used in the code.

def get_augmentation_model():
    """Build the data augmentation model."""
    data_augmentation = keras.Sequential(
        [
            layers.Resizing(..),
            layers.RandomCrop(...), # <--- possible causes
            layers.RandomFlip("horizontal"),
            layers.Rescaling(1 / 255.0),
        ]
    )
    return data_augmentation

Called a function referencing variables that have been deleted. This likely means that function-local variables were created and not referenced elsewhere in the program. This is generally a mistake; consider storing variables in an object attribute on the first call.

Though I looked over the source code but didn't feel any suspicious bug. So, to end this, I've implemented custom random crop layers which theoretically do the same. However, if you insist to use built-in layers.RandomCrop, then I would recommend opening a ticket on GitHub regarding this issue.

---

To make the whole code end-to-end training and saving+reloading, change the code as follows. First, implement the call method in the ShiftViTModel class.

class ShiftViTModel(keras.Model):
    def __init__(self):
        super().__init__(**kwargs)
        ...
        # init params

    def get_config(self):
        ...
        return config

    def _calculate_loss(self, data, training=False):
        ...
        return total_loss, labels, logits

    def train_step(self, inputs):
        ...
        return {m.name: m.result() for m in self.metrics}

    def test_step(self, data):
        ...
        return {m.name: m.result() for m in self.metrics}
    
    # implement the call function
    def call(self, images):
        augmented_images = self.data_augmentation(images)
        x = self.patch_projection(augmented_images)
        logits = self.global_avg_pool(x)
        return logits

Next, implement a custom random crop layer and use it as follows instead of layers.RandomCrop.

class CustomRandomCrop(layers.Layer):
    def __init__(self, size, **kwargs):
        super().__init__(**kwargs)
        self.size = size
        
    def call(self, inputs, training=True):
        if training:
            outputs = tf.map_fn(lambda img: tf.image.random_crop(img,
                                               self.size), inputs)
        else:
            outputs = tf.image.resize(inputs, self.size[:-1])
        return outputs
            
    def get_config(self):
        config = super().get_config()
        config.update(
            {
                'size': self.size,
            }
        )
        return config

def get_augmentation_model():
    """Build the data augmentation model."""
    data_augmentation = keras.Sequential(
        [
            layers.Resizing(...),
            CustomRandomCrop(...), # custom random crop
            layers.RandomFlip("horizontal"),
            layers.Rescaling(1 / 255.0),
        ]
    )
    return data_augmentation

After these changes, we can now do as follows without any error.

x,y = next(iter(train_ds))
print(x.shape, y.shape)
model(x.shape) 
# OK

history = model.fit(
    x=x, y=y,
    epochs=1
) 
# OK

model.evaluate(x, y)
# OK

Saving and reloading work as well.

model.save('/content/VIT-SHIFT')
# OK

new_model = tf.keras.models.load_model('/content/VIT-SHIFT')
# OK

np.testing.assert_allclose(
    model.predict(x), new_model.predict(x)
)
# OK

HERE is the complete working Code-in-Colab. Please save the file, I might erase the file from the drive someday.

Lastly, FYI, history.save('...'), you can't do that in keras. In order to save the tensorflow/keras model, please refer to this document. The history object will return only the tracked metrics and loss during training time. For example

history = model.fit(
    x=x, y=y,
    epochs=1
)

history.history

{'accuracy': [0.09765625],
 'loss': [6.204378128051758],
 'top-5-accuracy': [0.36328125]}

You can save the training logs from the above dictionary or better use the CSVLogger during training.