title gradient centralization for better training performance rishit daglihttps github comrishitdagli converted to keras 3 by muhammad anas razahttps anasrz com date created 061821 last modified 072523 description implement gradient centralization to improve training performance of dnns accelerator gpu introduction this example implements gradient centralizationhttps arxiv orgabs2004 01461 a new optimization technique for deep neural networks by yong et al and demonstrates it on laurence moroney s horses or humans datasethttps www tensorflow orgdatasetscataloghorsesorhumans gradient centralization can both speedup training process and improve the final generalization performance of dnns it operates directly on gradients by centralizing the gradient vectors to have zero mean gradient centralization morever improves the lipschitzness of the loss function and its gradient so that the training process becomes more efficient and stable this example requires tensorflowdatasets which can be installed with this command pip install tensorflowdatasets setup prepare the data for this example we will be using the horses or humans datasethttps www tensorflow orgdatasetscataloghorsesorhumans use data augmentation we will rescale the data to 0 1 and perform simple augmentations to our data helper to apply augmentation rescale dataset batch dataset use data augmentation only on the training set use buffered prefecting rescale and augment the data define a model in this section we will define a convolutional neural network implement gradient centralization we will now subclass the rmsprop optimizer class modifying the keras optimizers optimizer getgradients method where we now implement gradient centralization on a high level the idea is that let us say we obtain our gradients through back propogation for a dense or convolution layer we then compute the mean of the column vectors of the weight matrix and then remove the mean from each column vector the experiments in this paperhttps arxiv orgabs2004 01461 on various applications including general image classification finegrained image classification detection and segmentation and person reid demonstrate that gc can consistently improve the performance of dnn learning also for simplicity at the moment we are not implementing gradient cliiping functionality however this quite easy to implement at the moment we are just creating a subclass for the rmsprop optimizer however you could easily reproduce this for any other optimizer or on a custom optimizer in the same way we will be using this class in the later section when we train a model with gradient centralization we here just provide a modified getgradients function since we are trying to just compute the centralized gradients training utilities we will also create a callback which allows us to easily measure the total training time and the time taken for each epoch since we are interested in comparing the effect of gradient centralization on the model we built above train the model without gc we now train the model we built earlier without gradient centralization which we can compare to the training performance of the model trained with gradient centralization we also save the history since we later want to compare our model trained with and not trained with gradient centralization train the model with gc we will now train the same model this time using gradient centralization notice our optimizer is the one using gradient centralization this time comparing performance readers are encouraged to try out gradient centralization on different datasets from different domains and experiment with it s effect you are strongly advised to check out the original paperhttps arxiv orgabs2004 01461 as well the s present several studies on gradient centralization showing how it can improve general performance generalization training time as well as more efficient many thanks to ali mustufa shaikhhttps github comialimustufa for reviewing this implementation introduction this example implements gradient centralization https arxiv org abs 2004 01461 a new optimization technique for deep neural networks by yong et al and demonstrates it on laurence moroney s horses or humans dataset https www tensorflow org datasets catalog horses_or_humans gradient centralization can both speedup training process and improve the final generalization performance of dnns it operates directly on gradients by centralizing the gradient vectors to have zero mean gradient centralization morever improves the lipschitzness of the loss function and its gradient so that the training process becomes more efficient and stable this example requires tensorflow_datasets which can be installed with this command pip install tensorflow datasets setup prepare the data for this example we will be using the horses or humans dataset https www tensorflow org datasets catalog horses_or_humans use data augmentation we will rescale the data to 0 1 and perform simple augmentations to our data helper to apply augmentation rescale dataset batch dataset use data augmentation only on the training set use buffered prefecting rescale and augment the data define a model in this section we will define a convolutional neural network implement gradient centralization we will now subclass the rmsprop optimizer class modifying the keras optimizers optimizer get_gradients method where we now implement gradient centralization on a high level the idea is that let us say we obtain our gradients through back propogation for a dense or convolution layer we then compute the mean of the column vectors of the weight matrix and then remove the mean from each column vector the experiments in this paper https arxiv org abs 2004 01461 on various applications including general image classification fine grained image classification detection and segmentation and person reid demonstrate that gc can consistently improve the performance of dnn learning also for simplicity at the moment we are not implementing gradient cliiping functionality however this quite easy to implement at the moment we are just creating a subclass for the rmsprop optimizer however you could easily reproduce this for any other optimizer or on a custom optimizer in the same way we will be using this class in the later section when we train a model with gradient centralization we here just provide a modified get_gradients function since we are trying to just compute the centralized gradients training utilities we will also create a callback which allows us to easily measure the total training time and the time taken for each epoch since we are interested in comparing the effect of gradient centralization on the model we built above train the model without gc we now train the model we built earlier without gradient centralization which we can compare to the training performance of the model trained with gradient centralization we also save the history since we later want to compare our model trained with and not trained with gradient centralization train the model with gc we will now train the same model this time using gradient centralization notice our optimizer is the one using gradient centralization this time comparing performance readers are encouraged to try out gradient centralization on different datasets from different domains and experiment with it s effect you are strongly advised to check out the original paper https arxiv org abs 2004 01461 as well the s present several studies on gradient centralization showing how it can improve general performance generalization training time as well as more efficient many thanks to ali mustufa shaikh https github com ialimustufa for reviewing this implementation
from time import time

import keras
from keras import layers
from keras.optimizers import RMSprop
from keras import ops

from tensorflow import data as tf_data
import tensorflow_datasets as tfds




num_classes = 2
input_shape = (300, 300, 3)
dataset_name = "horses_or_humans"
batch_size = 128
AUTOTUNE = tf_data.AUTOTUNE

(train_ds, test_ds), metadata = tfds.load(
    name=dataset_name,
    split=[tfds.Split.TRAIN, tfds.Split.TEST],
    with_info=True,
    as_supervised=True,
)

print(f"Image shape: {metadata.features['image'].shape}")
print(f"Training images: {metadata.splits['train'].num_examples}")
print(f"Test images: {metadata.splits['test'].num_examples}")



rescale = layers.Rescaling(1.0 / 255)

data_augmentation = [
    layers.RandomFlip("horizontal_and_vertical"),
    layers.RandomRotation(0.3),
    layers.RandomZoom(0.2),
]



def apply_aug(x):
    for aug in data_augmentation:
        x = aug(x)
    return x


def prepare(ds, shuffle=False, augment=False):
    
    ds = ds.map(lambda x, y: (rescale(x), y), num_parallel_calls=AUTOTUNE)

    if shuffle:
        ds = ds.shuffle(1024)

    
    ds = ds.batch(batch_size)

    
    if augment:
        ds = ds.map(
            lambda x, y: (apply_aug(x), y),
            num_parallel_calls=AUTOTUNE,
        )

    
    return ds.prefetch(buffer_size=AUTOTUNE)




train_ds = prepare(train_ds, shuffle=True, augment=True)
test_ds = prepare(test_ds)


model = keras.Sequential(
    [
        layers.Conv2D(16, (3, 3), activation="relu", input_shape=(300, 300, 3)),
        layers.MaxPooling2D(2, 2),
        layers.Conv2D(32, (3, 3), activation="relu"),
        layers.Dropout(0.5),
        layers.MaxPooling2D(2, 2),
        layers.Conv2D(64, (3, 3), activation="relu"),
        layers.Dropout(0.5),
        layers.MaxPooling2D(2, 2),
        layers.Conv2D(64, (3, 3), activation="relu"),
        layers.MaxPooling2D(2, 2),
        layers.Conv2D(64, (3, 3), activation="relu"),
        layers.MaxPooling2D(2, 2),
        layers.Flatten(),
        layers.Dropout(0.5),
        layers.Dense(512, activation="relu"),
        layers.Dense(1, activation="sigmoid"),
    ]
)




class GCRMSprop(RMSprop):
    def get_gradients(self, loss, params):
        
        

        grads = []
        gradients = super().get_gradients()
        for grad in gradients:
            grad_len = len(grad.shape)
            if grad_len > 1:
                axis = list(range(grad_len - 1))
                grad -= ops.mean(grad, axis=axis, keep_dims=True)
            grads.append(grad)

        return grads


optimizer = GCRMSprop(learning_rate=1e-4)




class TimeHistory(keras.callbacks.Callback):
    def on_train_begin(self, logs={}):
        self.times = []

    def on_epoch_begin(self, batch, logs={}):
        self.epoch_time_start = time()

    def on_epoch_end(self, batch, logs={}):
        self.times.append(time() - self.epoch_time_start)




time_callback_no_gc = TimeHistory()
model.compile(
    loss="binary_crossentropy",
    optimizer=RMSprop(learning_rate=1e-4),
    metrics=["accuracy"],
)

model.summary()



history_no_gc = model.fit(
    train_ds, epochs=10, verbose=1, callbacks=[time_callback_no_gc]
)



time_callback_gc = TimeHistory()
model.compile(
    loss="binary_crossentropy", optimizer=optimizer, metrics=["accuracy"]
)

model.summary()

history_gc = model.fit(
    train_ds, epochs=10, verbose=1, callbacks=[time_callback_gc]
)



print("Not using Gradient Centralization")
print(f"Loss: {history_no_gc.history['loss'][-1]}")
print(f"Accuracy: {history_no_gc.history['accuracy'][-1]}")
print(f"Training Time: {sum(time_callback_no_gc.times)}")

print("Using Gradient Centralization")
print(f"Loss: {history_gc.history['loss'][-1]}")
print(f"Accuracy: {history_gc.history['accuracy'][-1]}")
print(f"Training Time: {sum(time_callback_gc.times)}")