coding utf8 lukewoodhttps lukewood xyz date created 03282023 last modified 07252023 description use kerascv to train powerful image classifiers introduction classification is the process of predicting a categorical label for a given input image while classification is a relatively straightforward computer vision task modern approaches still are built of several complex components luckily kerascv provides apis to construct commonly used components this guide demonstrates kerascv s modular approach to solving image classification problems at three levels of complexity inference with a pretrained classifier finetuning a pretrained backbone training a image classifier from scratch multibackend support kerascv s imageclassifier model supports several backends like jax pytorch and tensorflow with the help of keras to enable multibackend support in kerascv set the kerascvmultibackend environment variable we can then switch between different backends by setting the kerasbackend environment variable currently tensorflow jax and torch are supported this demonstration uses the jax backend inference with a pretrained classifier let s get started with the simplest kerascv api a pretrained classifier in this example we will construct a classifier that was pretrained on the imagenet dataset we ll use this model to solve the age old cat or dog problem the highest level module in kerascv is a task a task is a keras model consisting of a generally pretrained backbone model and taskspecific layers here s an example using kerascv models imageclassifier with an efficientnetv2b0 backbone efficientnetv2b0 is a great starting model when constructing an image classification pipeline this architecture manages to achieve high accuracy while using a parameter count of 7m if an efficientnetv2b0 is not powerful enough for the task you are hoping to solve be sure to check out kerascv s other available backboneshttps github comkerasteamkerascvtreemasterkerascvmodelsbackbones you may notice a small deviation from the old keras applications api where you would construct the class with efficientnetv2b0weightsimagenet while the old api was great for classification it did not scale effectively to other use cases that required complex architectures like object deteciton and semantic segmentation now that our classifier is built let s apply it to this cute cat picture next let s get some predictions from our classifier predictions classifier predictnp expanddimsimage axis0 topclasses predictions0 argsortaxis1 classes keras utils getfile originhttps gist githubusercontent comlukewood62eebcd5c5c4a4d0e0b7845780f76d55rawfde63e5e4c09e2fa0a3436680f436bdcb8325aacimagenetclassnames json with openclasses rb as f classes json loadf great both of these appear to be correct however one of the classes is velvet we re trying to classify cats vs dogs we don t care about the velvet blanket ideally we d have a classifier that only performs computation to determine if an image is a cat or a dog and has all of its resources dedicated to this task this can be solved by fine tuning our own classifier fine tuning a pretrained classifier when labeled images specific to our task are available finetuning a custom classifier can improve performance if we want to train a cats vs dogs classifier using explicitly labeled cat vs dog data should perform better than the generic classifier for many tasks no relevant pretrained model will be available e g categorizing images specific to your application first let s get started by loading some data staticly resize images as we only iterate the dataset once shuffle the dataset to increase diversity of batches 10batchsize follows the assumption that bigger machines can handle bigger shuffle buffers meow next let s construct our model the use of imagenet in the preset name indicates that the backbone was pretrained on the imagenet dataset pretrained backbones extract more information from our labeled examples by leveraging patterns extracted from potentially much larger datasets next lets put together our classifier here our classifier is just a simple keras sequential all that is left to do is call model fit let s look at how our model performs after the fine tuning predictions model predictnp expanddimsimage axis0 classes 0 cat 1 dog printtop class is classespredictions0 argmax train a classifier from scratch numclasses 101 change epochs to 100 to fully train epochs 1 encoder keras layers categoryencodingnumclasses onehot dtypeint32 def packageinputsimage label return images image labels encoderlabel trainds evalds tfds load caltech101 splittrain test assupervisedtrue trainds trainds mappackageinputs numparallelcallstfdata autotune evalds evalds mappackageinputs numparallelcallstfdata autotune trainds trainds shufflebatchsize 16 trainds trainds raggedbatchbatchsize evalds evalds raggedbatchbatchsize batch nextitertrainds take1 imagebatch batchimages labelbatch batchlabels kerascv visualization plotimagegallery imagebatch totensor rows3 cols3 valuerange0 255 showtrue randomflip kerascv layers randomflip augmenters randomflip imagebatch randomflipimagebatch kerascv visualization plotimagegallery imagebatch totensor rows3 cols3 valuerange0 255 showtrue cropandresize kerascv layers randomcropandresize targetsizeimagesize cropareafactor0 8 1 0 aspectratiofactor0 9 1 1 augmenters cropandresize imagebatch cropandresizeimagebatch kerascv visualization plotimagegallery imagebatch rows3 cols3 valuerange0 255 showtrue randaugment kerascv layers randaugment augmentationsperimage3 magnitude0 3 valuerange0 255 augmenters randaugment imagebatch randaugmentimagebatch kerascv visualization plotimagegallery imagebatch rows3 cols3 valuerange0 255 showtrue randomcutout kerascv layers randomcutout widthfactor0 4 heightfactor0 4 kerascv visualization plotimagegallery randomcutoutimagebatch rows3 cols3 valuerange0 255 showtrue cutmix kerascv layers cutmix cutmix needs to modify both images and labels inputs images imagebatch labels tf castlabelbatch float32 kerascv visualization plotimagegallery cutmixinputsimages rows3 cols3 valuerange0 255 showtrue mixup kerascv layers mixup mixup needs to modify both images and labels inputs images imagebatch labels tf castlabelbatch float32 kerascv visualization plotimagegallery mixupinputsimages rows3 cols3 valuerange0 255 showtrue cutmixormixup kerascv layers randomchoice cutmix mixup batchwisetrue augmenters cutmixormixup we also need to resize our evaluation set to get dense batches of the image size expected by our model we use the deterministic kerascv layers resizing in this case to avoid adding noise to our evaluation metric finally lets unpackage our datasets and prepare to pass them to model fit which accepts a tuple of images labels data augmentation is by far the hardest piece of training a modern classifier congratulations on making it this far optimizer tuning to achieve optimal performance we need to use a learning rate schedule instead of a single learning rate while we won t go into detail on the cosine decay with warmup schedule used here you can read more about it herehttps scorrea92 medium comcosinelearningratedecaye8b50aa455b cosine decay warmupcosinedecay schedulehttps i imgur comycr5pii png the schedule looks a as we expect next let s construct this optimizer at long last we can now build our model and call fit kerascv models efficientnetv2b0backbone is a convenience alias for kerascv models efficientnetv2backbone frompreset efficientnetv2b0 note that this preset does not come with any pretrained weights since the labels produced by mixup and cutmix are somewhat artificial we employ label smoothing to prevent the model from overfitting to artifacts of this augmentation process let s compile our model model compile lossloss optimizeroptimizer metrics metrics categoricalaccuracy metrics topkcategoricalaccuracyk5 congratulations you now know how to train a powerful image classifier from scratch in kerascv depending on the availability of labeled data for your application training from scratch may or may not be more powerful than using transfer learning in addition to the data augmentations discussed above for smaller datasets pretrained models generally produce high accuracy and faster convergence conclusions while image classification is perhaps the simplest problem in computer vision the modern landscape has numerous complex components luckily kerascv offers robust productiongrade apis to make assembling most of these components possible in one line of code through the use of kerascv s imageclassifier api pretrained weights and kerascv data augmentations you can assemble everything you need to train a powerful classifier in a few hundred lines of code as a follow up exercise give the following a try fine tune a kerascv classifier on your own dataset learn more about kerascv s data augmentationshttps keras ioguideskerascvcutmixmixupandrandaugment check out how we train our models on imagenethttps github comkerasteamkerascvblobmasterexamplestrainingclassificationimagenetbasictraining py coding utf 8 lukewood https lukewood xyz date created 03 28 2023 last modified 07 25 2023 description use kerascv to train powerful image classifiers introduction classification is the process of predicting a categorical label for a given input image while classification is a relatively straightforward computer vision task modern approaches still are built of several complex components luckily kerascv provides apis to construct commonly used components this guide demonstrates kerascv s modular approach to solving image classification problems at three levels of complexity inference with a pretrained classifier fine tuning a pretrained backbone training a image classifier from scratch multi backend support kerascv s imageclassifier model supports several backends like jax pytorch and tensorflow with the help of keras to enable multi backend support in kerascv set the keras_cv_multi_backend environment variable we can then switch between different backends by setting the keras_backend environment variable currently tensorflow jax and torch are supported this demonstration uses the jax backend inference with a pretrained classifier let s get started with the simplest kerascv api a pretrained classifier in this example we will construct a classifier that was pretrained on the imagenet dataset we ll use this model to solve the age old cat or dog problem the highest level module in kerascv is a task a task is a keras model consisting of a generally pretrained backbone model and task specific layers here s an example using keras_cv models imageclassifier with an efficientnetv2b0 backbone efficientnetv2b0 is a great starting model when constructing an image classification pipeline this architecture manages to achieve high accuracy while using a parameter count of 7m if an efficientnetv2b0 is not powerful enough for the task you are hoping to solve be sure to check out kerascv s other available backbones https github com keras team keras cv tree master keras_cv models backbones you may notice a small deviation from the old keras applications api where you would construct the class with efficientnetv2b0 weights imagenet while the old api was great for classification it did not scale effectively to other use cases that required complex architectures like object deteciton and semantic segmentation now that our classifier is built let s apply it to this cute cat picture next let s get some predictions from our classifier predictions come in the form of softmax ed category rankings we can find the index of the top classes using a simple argsort function in order to decode the class mappings we can construct a mapping from category indices to imagenet class names for convenience i ve stored the imagenet class mapping in a github gist let s download and load it now now we can simply look up the class names via index great both of these appear to be correct however one of the classes is velvet we re trying to classify cats vs dogs we don t care about the velvet blanket ideally we d have a classifier that only performs computation to determine if an image is a cat or a dog and has all of its resources dedicated to this task this can be solved by fine tuning our own classifier fine tuning a pretrained classifier when labeled images specific to our task are available fine tuning a custom classifier can improve performance if we want to train a cats vs dogs classifier using explicitly labeled cat vs dog data should perform better than the generic classifier for many tasks no relevant pretrained model will be available e g categorizing images specific to your application first let s get started by loading some data staticly resize images as we only iterate the dataset once shuffle the dataset to increase diversity of batches 10 batch_size follows the assumption that bigger machines can handle bigger shuffle buffers meow next let s construct our model the use of imagenet in the preset name indicates that the backbone was pretrained on the imagenet dataset pretrained backbones extract more information from our labeled examples by leveraging patterns extracted from potentially much larger datasets next lets put together our classifier here our classifier is just a simple keras sequential all that is left to do is call model fit let s look at how our model performs after the fine tuning awesome looks like the model correctly classified the image train a classifier from scratch now that we ve gotten our hands dirty with classification let s take on one last task training a classification model from scratch a standard benchmark for image classification is the imagenet dataset however due to licensing constraints we will use the caltech 101 image classification dataset in this tutorial while we use the simpler caltech 101 dataset in this guide the same training template may be used on imagenet to achieve near state of the art scores let s start out by tackling data loading change epochs to 100 to fully train the caltech101 dataset has different sizes for every image so we use the ragged_batch api to batch them together while maintaining each individual image s shape information data augmentation in our previous finetuning exmaple we performed a static resizing operation and did not utilize any image augmentation this is because a single pass over the training set was sufficient to achieve decent results when training to solve a more difficult task you ll want to include data augmentation in your data pipeline data augmentation is a technique to make your model robust to changes in input data such as lighting cropping and orientation kerascv includes some of the most useful augmentations in the keras_cv layers api creating an optimal pipeline of augmentations is an art but in this section of the guide we ll offer some tips on best practices for classification one caveat to be aware of with image data augmentation is that you must be careful to not shift your augmented data distribution too far from the original data distribution the goal is to prevent overfitting and increase generalization but samples that lie completely out of the data distribution simply add noise to the training process the first augmentation we ll use is randomflip this augmentation behaves more or less how you d expect it either flips the image or not while this augmentation is useful in caltech101 and imagenet it should be noted that it should not be used on tasks where the data distribution is not vertical mirror invariant an example of a dataset where this occurs is mnist hand written digits flipping a 6 over the vertical axis will make the digit appear more like a 7 than a 6 but the label will still show a 6 half of the images have been flipped the next augmentation we ll use is randomcropandresize this operation selects a random subset of the image then resizes it to the provided target size by using this augmentation we force our classifier to become spatially invariant additionally this layer accepts an aspect_ratio_factor which can be used to distort the aspect ratio of the image while this can improve model performance it should be used with caution it is very easy for an aspect ratio distortion to shift a sample too far from the original training set s data distribution remember the goal of data augmentation is to produce more training samples that align with the data distribution of your training set randomcropandresize also can handle tf raggedtensor inputs in the caltech101 image dataset images come in a wide variety of sizes as such they cannot easily be batched together into a dense training batch luckily randomcropandresize handles the ragged dense conversion process for you let s add a randomcropandresize to our set of augmentations great we are now working with a batch of dense images next up lets include some spatial and color based jitter to our training set this will allow us to produce a classifier that is robust to lighting flickers shadows and more there are limitless ways to augment an image by altering color and spatial features but perhaps the most battle tested technique is randaugment https arxiv org abs 1909 13719 randaugment is actually a set of 10 different augmentations autocontrast equalize solarize randomcolorjitter randomcontrast randombrightness shearx sheary translatex and translatey at inference time num_augmentations augmenters are sampled for each image and random magnitude factors are sampled for each these augmentations are then applied sequentially kerascv makes tuning these parameters easy using the augmentations_per_image and magnitude parameters let s take it for a spin looks great but we re not done yet what if an image is missing one critical feature of a class for example what if a leaf is blocking the view of a cat s ear but our classifier learned to classify cats simply by observing their ears one easy approach to tackling this is to use randomcutout which randomly strips out a sub section of the image while this tackles the problem reasonably well it can cause the classifier to develop responses to borders between features and black pixel areas caused by the cutout cutmix https arxiv org abs 1905 04899 solves the same issue by using a more complex and more effective technique instead of replacing the cut out areas with black pixels cutmix replaces these regions with regions of other images sampled from within your training set following this replacement the image s classification label is updated to be a blend of the original and mixed image s class label what does this look like in practice let s check it out cutmix needs to modify both images and labels let s hold off from adding it to our augmenter for a minute more on that soon next let s look into mixup unfortunately while mixup has been empirically shown to substantially improve both the robustness and the generalization of the trained model it is not well understood why such improvement occurs but a little alchemy never hurt anyone mixup works by sampling two images from a batch then proceeding to literally blend together their pixel intensities as well as their classification labels let s see it in action mixup needs to modify both images and labels if you look closely you ll see that the images have been blended together instead of applying cutmix and mixup to every image we instead pick one or the other to apply to each batch this can be expressed using keras_cv layers randomchoice now let s apply our final augmenter to the training data we also need to resize our evaluation set to get dense batches of the image size expected by our model we use the deterministic keras_cv layers resizing in this case to avoid adding noise to our evaluation metric finally lets unpackage our datasets and prepare to pass them to model fit which accepts a tuple of images labels data augmentation is by far the hardest piece of training a modern classifier congratulations on making it this far optimizer tuning to achieve optimal performance we need to use a learning rate schedule instead of a single learning rate while we won t go into detail on the cosine decay with warmup schedule used here you can read more about it here https scorrea92 medium com cosine learning rate decay e8b50aa455b cosine decay warmupcosinedecay schedule https i imgur com ycr5pii png the schedule looks a as we expect next let s construct this optimizer at long last we can now build our model and call fit keras_cv models efficientnetv2b0backbone is a convenience alias for keras_cv models efficientnetv2backbone from_preset efficientnetv2_b0 note that this preset does not come with any pretrained weights since the labels produced by mixup and cutmix are somewhat artificial we employ label smoothing to prevent the model from overfitting to artifacts of this augmentation process let s compile our model and finally call fit congratulations you now know how to train a powerful image classifier from scratch in kerascv depending on the availability of labeled data for your application training from scratch may or may not be more powerful than using transfer learning in addition to the data augmentations discussed above for smaller datasets pretrained models generally produce high accuracy and faster convergence conclusions while image classification is perhaps the simplest problem in computer vision the modern landscape has numerous complex components luckily kerascv offers robust production grade apis to make assembling most of these components possible in one line of code through the use of kerascv s imageclassifier api pretrained weights and kerascv data augmentations you can assemble everything you need to train a powerful classifier in a few hundred lines of code as a follow up exercise give the following a try fine tune a kerascv classifier on your own dataset learn more about kerascv s data augmentations https keras io guides keras_cv cut_mix_mix_up_and_rand_augment check out how we train our models on imagenet https github com keras team keras cv blob master examples training classification imagenet basic_training py
import os

os.environ["KERAS_CV_MULTI_BACKEND"] = "1"
os.environ["KERAS_BACKEND"] = "jax"

import json
import math
import keras_cv
import keras
from keras import ops
from keras import losses
from keras import optimizers
from keras.optimizers import schedules
from keras import metrics
import tensorflow as tf
from tensorflow import data as tf_data
import tensorflow_datasets as tfds
import numpy as np



classifier = keras_cv.models.ImageClassifier.from_preset(
    "efficientnetv2_b0_imagenet_classifier"
)



filepath = keras.utils.get_file(origin="https://i.imgur.com/9i63gLN.jpg")
image = keras.utils.load_img(filepath)
image = np.array(image)
keras_cv.visualization.plot_image_gallery(
    image[None, ...], rows=1, cols=1, value_range=(0, 255), show=True, scale=4
)



predictions = classifier.predict(np.expand_dims(image, axis=0))



top_classes = predictions[0].argsort(axis=-1)



classes = keras.utils.get_file(
    origin="https://gist.githubusercontent.com/LukeWood/62eebcd5c5c4a4d0e0b7845780f76d55/raw/fde63e5e4c09e2fa0a3436680f436bdcb8325aac/ImagenetClassnames.json"
)
with open(classes, "rb") as f:
    classes = json.load(f)



top_two = [classes[str(i)] for i in top_classes[-2:]]
print("Top two classes are:", top_two)



BATCH_SIZE = 32
IMAGE_SIZE = (224, 224)
AUTOTUNE = tf_data.AUTOTUNE
tfds.disable_progress_bar()

data, dataset_info = tfds.load(
    "cats_vs_dogs", with_info=True, as_supervised=True
)
train_steps_per_epoch = dataset_info.splits["train"].num_examples // BATCH_SIZE
train_dataset = data["train"]

num_classes = dataset_info.features["label"].num_classes

resizing = keras_cv.layers.Resizing(
    IMAGE_SIZE[0], IMAGE_SIZE[1], crop_to_aspect_ratio=True
)
encoder = keras.layers.CategoryEncoding(num_classes, "one_hot", dtype="int32")


def preprocess_inputs(image, label):
    
    return resizing(image), encoder(label)





train_dataset = train_dataset.shuffle(
    10 * BATCH_SIZE, reshuffle_each_iteration=True
).map(preprocess_inputs, num_parallel_calls=AUTOTUNE)
train_dataset = train_dataset.batch(BATCH_SIZE)

images = next(iter(train_dataset.take(1)))[0]
keras_cv.visualization.plot_image_gallery(images, value_range=(0, 255))



model = keras_cv.models.ImageClassifier.from_preset(
    "efficientnetv2_b0_imagenet", num_classes=2
)
model.compile(
    loss="categorical_crossentropy",
    optimizer=keras.optimizers.SGD(learning_rate=0.01),
    metrics=["accuracy"],
)



model.fit(train_dataset)



predictions = model.predict(np.expand_dims(image, axis=0))

classes = {0: "cat", 1: "dog"}
print("Top class is:", classes[predictions[0].argmax()])



NUM_CLASSES = 101

EPOCHS = 1

encoder = keras.layers.CategoryEncoding(NUM_CLASSES, "one_hot", dtype="int32")


def package_inputs(image, label):
    return {"images": image, "labels": encoder(label)}


train_ds, eval_ds = tfds.load(
    "caltech101", split=["train", "test"], as_supervised="true"
)
train_ds = train_ds.map(package_inputs, num_parallel_calls=tf_data.AUTOTUNE)
eval_ds = eval_ds.map(package_inputs, num_parallel_calls=tf_data.AUTOTUNE)

train_ds = train_ds.shuffle(BATCH_SIZE * 16)



train_ds = train_ds.ragged_batch(BATCH_SIZE)
eval_ds = eval_ds.ragged_batch(BATCH_SIZE)

batch = next(iter(train_ds.take(1)))
image_batch = batch["images"]
label_batch = batch["labels"]

keras_cv.visualization.plot_image_gallery(
    image_batch.to_tensor(),
    rows=3,
    cols=3,
    value_range=(0, 255),
    show=True,
)



random_flip = keras_cv.layers.RandomFlip()
augmenters = [random_flip]

image_batch = random_flip(image_batch)
keras_cv.visualization.plot_image_gallery(
    image_batch.to_tensor(),
    rows=3,
    cols=3,
    value_range=(0, 255),
    show=True,
)



crop_and_resize = keras_cv.layers.RandomCropAndResize(
    target_size=IMAGE_SIZE,
    crop_area_factor=(0.8, 1.0),
    aspect_ratio_factor=(0.9, 1.1),
)
augmenters += [crop_and_resize]

image_batch = crop_and_resize(image_batch)
keras_cv.visualization.plot_image_gallery(
    image_batch,
    rows=3,
    cols=3,
    value_range=(0, 255),
    show=True,
)



rand_augment = keras_cv.layers.RandAugment(
    augmentations_per_image=3,
    magnitude=0.3,
    value_range=(0, 255),
)
augmenters += [rand_augment]

image_batch = rand_augment(image_batch)
keras_cv.visualization.plot_image_gallery(
    image_batch,
    rows=3,
    cols=3,
    value_range=(0, 255),
    show=True,
)



random_cutout = keras_cv.layers.RandomCutout(
    width_factor=0.4, height_factor=0.4
)
keras_cv.visualization.plot_image_gallery(
    random_cutout(image_batch),
    rows=3,
    cols=3,
    value_range=(0, 255),
    show=True,
)



cut_mix = keras_cv.layers.CutMix()

inputs = {"images": image_batch, "labels": tf.cast(label_batch, "float32")}

keras_cv.visualization.plot_image_gallery(
    cut_mix(inputs)["images"],
    rows=3,
    cols=3,
    value_range=(0, 255),
    show=True,
)



mix_up = keras_cv.layers.MixUp()

inputs = {"images": image_batch, "labels": tf.cast(label_batch, "float32")}

keras_cv.visualization.plot_image_gallery(
    mix_up(inputs)["images"],
    rows=3,
    cols=3,
    value_range=(0, 255),
    show=True,
)



cut_mix_or_mix_up = keras_cv.layers.RandomChoice(
    [cut_mix, mix_up], batchwise=True
)
augmenters += [cut_mix_or_mix_up]



augmenter = keras_cv.layers.Augmenter(augmenters)
train_ds = train_ds.map(augmenter, num_parallel_calls=tf_data.AUTOTUNE)

image_batch = next(iter(train_ds.take(1)))["images"]
keras_cv.visualization.plot_image_gallery(
    image_batch,
    rows=3,
    cols=3,
    value_range=(0, 255),
    show=True,
)



inference_resizing = keras_cv.layers.Resizing(
    IMAGE_SIZE[0], IMAGE_SIZE[1], crop_to_aspect_ratio=True
)
eval_ds = eval_ds.map(inference_resizing, num_parallel_calls=tf_data.AUTOTUNE)

inference_resizing = keras_cv.layers.Resizing(
    IMAGE_SIZE[0], IMAGE_SIZE[1], crop_to_aspect_ratio=True
)
eval_ds = eval_ds.map(inference_resizing, num_parallel_calls=tf_data.AUTOTUNE)

image_batch = next(iter(eval_ds.take(1)))["images"]
keras_cv.visualization.plot_image_gallery(
    image_batch,
    rows=3,
    cols=3,
    value_range=(0, 255),
    show=True,
)




def unpackage_dict(inputs):
    return inputs["images"], inputs["labels"]


train_ds = train_ds.map(unpackage_dict, num_parallel_calls=tf_data.AUTOTUNE)
eval_ds = eval_ds.map(unpackage_dict, num_parallel_calls=tf_data.AUTOTUNE)




def lr_warmup_cosine_decay(
    global_step,
    warmup_steps,
    hold=0,
    total_steps=0,
    start_lr=0.0,
    target_lr=1e-2,
):
    
    learning_rate = (
        0.5
        * target_lr
        * (
            1
            + ops.cos(
                math.pi
                * ops.convert_to_tensor(
                    global_step - warmup_steps - hold, dtype="float32"
                )
                / ops.convert_to_tensor(
                    total_steps - warmup_steps - hold, dtype="float32"
                )
            )
        )
    )

    warmup_lr = target_lr * (global_step / warmup_steps)

    if hold > 0:
        learning_rate = ops.where(
            global_step > warmup_steps + hold, learning_rate, target_lr
        )

    learning_rate = ops.where(
        global_step < warmup_steps, warmup_lr, learning_rate
    )
    return learning_rate


class WarmUpCosineDecay(schedules.LearningRateSchedule):
    def __init__(
        self, warmup_steps, total_steps, hold, start_lr=0.0, target_lr=1e-2
    ):
        super().__init__()
        self.start_lr = start_lr
        self.target_lr = target_lr
        self.warmup_steps = warmup_steps
        self.total_steps = total_steps
        self.hold = hold

    def __call__(self, step):
        lr = lr_warmup_cosine_decay(
            global_step=step,
            total_steps=self.total_steps,
            warmup_steps=self.warmup_steps,
            start_lr=self.start_lr,
            target_lr=self.target_lr,
            hold=self.hold,
        )

        return ops.where(step > self.total_steps, 0.0, lr)




total_images = 9000
total_steps = (total_images // BATCH_SIZE) * EPOCHS
warmup_steps = int(0.1 * total_steps)
hold_steps = int(0.45 * total_steps)
schedule = WarmUpCosineDecay(
    start_lr=0.05,
    target_lr=1e-2,
    warmup_steps=warmup_steps,
    total_steps=total_steps,
    hold=hold_steps,
)
optimizer = optimizers.SGD(
    weight_decay=5e-4,
    learning_rate=schedule,
    momentum=0.9,
)



backbone = keras_cv.models.ResNet18V2Backbone()
model = keras.Sequential(
    [
        backbone,
        keras.layers.GlobalMaxPooling2D(),
        keras.layers.Dropout(rate=0.5),
        keras.layers.Dense(101, activation="softmax"),
    ]
)



loss = losses.CategoricalCrossentropy(label_smoothing=0.1)



model.compile(
    loss=loss,
    optimizer=optimizer,
    metrics=[
        metrics.CategoricalAccuracy(),
        metrics.TopKCategoricalAccuracy(k=5),
    ],
)



model.fit(
    train_ds,
    epochs=EPOCHS,
    validation_data=eval_ds,
)