title image classification using bigtransfer bit sayan nathhttps twitter comsayannath2350 date created 20210924 last modified 20210924 description bigtransfer bit stateoftheart transfer learning for image classification accelerator gpu introduction bigtransfer also known as bit is a stateoftheart transfer learning method for image classification transfer of pretrained representations improves sample efficiency and simplifies hyperparameter tuning when training deep neural networks for vision bit revisit the paradigm of pretraining on large supervised datasets and finetuning the model on a target task the importance of appropriately choosing normalization layers and scaling the architecture capacity as the amount of pretraining data increases bigtransferbit is trained on public datasets along with code in tf2 jax and pytorchhttps github comgoogleresearchbigtransfer this will help anyone to reach state of the art performance on their task of interest even with just a handful of labeled images per class you can find bit models pretrained on imagenethttps imagenet orgchallengeslsvrc2012index and imagenet21k in tfhubhttps tfhub devgooglecollectionsbit1 as tensorflow2 savedmodels that you can use easily as keras layers there are a variety of sizes ranging from a standard resnet50 to a resnet152x4 152 layers deep 4x wider than a typical resnet50 for users with larger computational and memory budgets but higher accuracy requirements https i imgur comxewvfe7 jpeg figure the xaxis shows the number of images used per class ranging from 1 to the full dataset on the plots on the left the curve in blue above is our bitl model whereas the curve below is a resnet50 pretrained on imagenet ilsvrc2012 setup gather flower dataset visualise the dataset define hyperparameters the hyperparamteres like schedulelength and scheduleboundaries are determined based on empirical results the method has been explained in the original paperhttps arxiv orgabs1912 11370 and in their google ai blog posthttps ai googleblog com202005opensourcingbitexploringlargescale html the schedulelength is aslo determined whether to use mixup augmentationhttps arxiv orgabs1710 09412 or not you can also find an easy mixup implementation in keras coding exampleshttps keras ioexamplesvisionmixup https i imgur comosaibyz jpeg define preprocessing helper functions define the data pipeline training pipeline validation pipeline visualise the training samples load pretrained tfhub model into a keraslayer create bigtransfer bit model to create the new model we 1 cut off the bit models original head this leaves us with the prelogits output we do not have to do this if we use the feature extractor models i e all those in subdirectories titled featurevectors since for those models the head has already been cut off 2 add a new head with the number of outputs equal to the number of classes of our new task note that it is important that we initialise the head to all zeroes define optimizer and loss decay learning rate by a factor of 10 at scheduleboundaries compile the model set up callbacks train the model plot the training and validation metrics evaluate the model conclusion bit performs well across a surprisingly wide range of data regimes from 1 example per class to 1m total examples bit achieves 87 5 top1 accuracy on ilsvrc2012 99 4 on cifar10 and 76 3 on the 19 task visual task adaptation benchmark vtab on small datasets bit attains 76 8 on ilsvrc2012 with 10 examples per class and 97 0 on cifar10 with 10 examples per class https i imgur comb1lw5fz png you can experiment further with the bigtransfer method by following the original paperhttps arxiv orgabs1912 11370 example available on huggingface trained model demo generic badgehttps img shields iobadgef09fa49720modelbitblack svghttps huggingface cokerasiobit generic badgehttps img shields iobadgef09fa49720spacesbitblack svghttps huggingface cospaceskerasiosiamesecontrastive introduction bigtransfer also known as bit is a state of the art transfer learning method for image classification transfer of pre trained representations improves sample efficiency and simplifies hyperparameter tuning when training deep neural networks for vision bit revisit the paradigm of pre training on large supervised datasets and fine tuning the model on a target task the importance of appropriately choosing normalization layers and scaling the architecture capacity as the amount of pre training data increases bigtransfer bit is trained on public datasets along with code in tf2 jax and pytorch https github com google research big_transfer this will help anyone to reach state of the art performance on their task of interest even with just a handful of labeled images per class you can find bit models pre trained on imagenet https image net org challenges lsvrc 2012 index and imagenet 21k in tfhub https tfhub dev google collections bit 1 as tensorflow2 savedmodels that you can use easily as keras layers there are a variety of sizes ranging from a standard resnet50 to a resnet152x4 152 layers deep 4x wider than a typical resnet50 for users with larger computational and memory budgets but higher accuracy requirements https i imgur com xewvfe7 jpeg figure the x axis shows the number of images used per class ranging from 1 to the full dataset on the plots on the left the curve in blue above is our bit l model whereas the curve below is a resnet 50 pre trained on imagenet ilsvrc 2012 setup gather flower dataset visualise the dataset define hyperparameters optimise the pipeline performance number of classes we will train on lower resolution images and will still attain good results more the dataset size the schedule length increase the hyperparamteres like schedule_length and schedule_boundaries are determined based on empirical results the method has been explained in the original paper https arxiv org abs 1912 11370 and in their google ai blog post https ai googleblog com 2020 05 open sourcing bit exploring large scale html the schedule_length is aslo determined whether to use mixup augmentation https arxiv org abs 1710 09412 or not you can also find an easy mixup implementation in keras coding examples https keras io examples vision mixup https i imgur com osaibyz jpeg define preprocessing helper functions to ensure at least there are 50 epochs of training define the data pipeline training pipeline repeat dataset_size num_steps validation pipeline visualise the training samples load pretrained tf hub model into a keraslayer create bigtransfer bit model to create the new model we 1 cut off the bit model s original head this leaves us with the pre logits output we do not have to do this if we use the feature extractor models i e all those in subdirectories titled feature_vectors since for those models the head has already been cut off 2 add a new head with the number of outputs equal to the number of classes of our new task note that it is important that we initialise the head to all zeroes define optimizer and loss decay learning rate by a factor of 10 at schedule_boundaries compile the model set up callbacks train the model plot the training and validation metrics evaluate the model conclusion bit performs well across a surprisingly wide range of data regimes from 1 example per class to 1m total examples bit achieves 87 5 top 1 accuracy on ilsvrc 2012 99 4 on cifar 10 and 76 3 on the 19 task visual task adaptation benchmark vtab on small datasets bit attains 76 8 on ilsvrc 2012 with 10 examples per class and 97 0 on cifar 10 with 10 examples per class https i imgur com b1lw5fz png you can experiment further with the bigtransfer method by following the original paper https arxiv org abs 1912 11370 example available on huggingface trained model demo generic badge https img shields io badge f0 9f a4 97 20model bit black svg https huggingface co keras io bit generic badge https img shields io badge f0 9f a4 97 20spaces bit black svg https huggingface co spaces keras io siamese contrastive
import os

os.environ["KERAS_BACKEND"] = "tensorflow"

import keras
import numpy as np
import matplotlib.pyplot as plt

import tensorflow as tf
import tensorflow_hub as hub
import tensorflow_datasets as tfds

tfds.disable_progress_bar()

SEEDS = 42

np.random.seed(SEEDS)
keras.utils.set_random_seed(SEEDS)



train_ds, validation_ds = tfds.load(
    "tf_flowers",
    split=["train[:85%]", "train[85%:]"],
    as_supervised=True,
)



plt.figure(figsize=(10, 10))
for i, (image, label) in enumerate(train_ds.take(9)):
    ax = plt.subplot(3, 3, i + 1)
    plt.imshow(image)
    plt.title(int(label))
    plt.axis("off")



RESIZE_TO = 384
CROP_TO = 224
BATCH_SIZE = 64
STEPS_PER_EPOCH = 10
AUTO = tf.data.AUTOTUNE  
NUM_CLASSES = 5  
SCHEDULE_LENGTH = 500  
SCHEDULE_BOUNDARIES = [
    200,
    300,
    400,
]  





SCHEDULE_LENGTH = SCHEDULE_LENGTH * 512 / BATCH_SIZE


@tf.function
def preprocess_train(image, label):
    image = tf.image.random_flip_left_right(image)
    image = tf.image.resize(image, (RESIZE_TO, RESIZE_TO))
    image = tf.image.random_crop(image, (CROP_TO, CROP_TO, 3))
    image = image / 255.0
    return (image, label)


@tf.function
def preprocess_test(image, label):
    image = tf.image.resize(image, (RESIZE_TO, RESIZE_TO))
    image = image / 255.0
    return (image, label)


DATASET_NUM_TRAIN_EXAMPLES = train_ds.cardinality().numpy()

repeat_count = int(
    SCHEDULE_LENGTH * BATCH_SIZE / DATASET_NUM_TRAIN_EXAMPLES * STEPS_PER_EPOCH
)
repeat_count += 50 + 1  




pipeline_train = (
    train_ds.shuffle(10000)
    .repeat(repeat_count)  
    .map(preprocess_train, num_parallel_calls=AUTO)
    .batch(BATCH_SIZE)
    .prefetch(AUTO)
)


pipeline_validation = (
    validation_ds.map(preprocess_test, num_parallel_calls=AUTO)
    .batch(BATCH_SIZE)
    .prefetch(AUTO)
)



image_batch, label_batch = next(iter(pipeline_train))

plt.figure(figsize=(10, 10))
for n in range(25):
    ax = plt.subplot(5, 5, n + 1)
    plt.imshow(image_batch[n])
    plt.title(label_batch[n].numpy())
    plt.axis("off")



bit_model_url = "https://tfhub.dev/google/bit/m-r50x1/1"
bit_module = hub.KerasLayer(bit_model_url)




class MyBiTModel(keras.Model):
    def __init__(self, num_classes, module, **kwargs):
        super().__init__(**kwargs)

        self.num_classes = num_classes
        self.head = keras.layers.Dense(num_classes, kernel_initializer="zeros")
        self.bit_model = module

    def call(self, images):
        bit_embedding = self.bit_model(images)
        return self.head(bit_embedding)


model = MyBiTModel(num_classes=NUM_CLASSES, module=bit_module)



learning_rate = 0.003 * BATCH_SIZE / 512


lr_schedule = keras.optimizers.schedules.PiecewiseConstantDecay(
    boundaries=SCHEDULE_BOUNDARIES,
    values=[
        learning_rate,
        learning_rate * 0.1,
        learning_rate * 0.01,
        learning_rate * 0.001,
    ],
)
optimizer = keras.optimizers.SGD(learning_rate=lr_schedule, momentum=0.9)

loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True)



model.compile(optimizer=optimizer, loss=loss_fn, metrics=["accuracy"])



train_callbacks = [
    keras.callbacks.EarlyStopping(
        monitor="val_accuracy", patience=2, restore_best_weights=True
    )
]



history = model.fit(
    pipeline_train,
    batch_size=BATCH_SIZE,
    epochs=int(SCHEDULE_LENGTH / STEPS_PER_EPOCH),
    steps_per_epoch=STEPS_PER_EPOCH,
    validation_data=pipeline_validation,
    callbacks=train_callbacks,
)




def plot_hist(hist):
    plt.plot(hist.history["accuracy"])
    plt.plot(hist.history["val_accuracy"])
    plt.plot(hist.history["loss"])
    plt.plot(hist.history["val_loss"])
    plt.title("Training Progress")
    plt.ylabel("Accuracy/Loss")
    plt.xlabel("Epochs")
    plt.legend(
        ["train_acc", "val_acc", "train_loss", "val_loss"], loc="upper left"
    )
    plt.show()


plot_hist(history)



accuracy = model.evaluate(pipeline_validation)[1] * 100
print("Accuracy: {:.2f}%".format(accuracy))