Add BigTransfer as TFHub.KerasLayer

[avber]

Hi @leondgarse ,

are there any obstacles to adding BigTransfer as TFHub.KerasLayer. No need to copy/reimplement their code.

Though it is important to initialise model heads to all zeroes if BIT is used as a backbone.

Just a few lines from their example

import tensorflow_hub as hub

Load model from TFHub into KerasLayer

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

class MyBiTModel(tf.keras.Model):
"""BiT with a new head."""

def init(self, num_classes, module):
super().init()

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

def call(self, images):
# No need to cut head off since we are using feature extractor model
bit_embedding = self.bit_model(images)
return self.head(bit_embedding)

model = MyBiTModel(num_classes=5, module=module)

[leondgarse]

I'm not very familliar with TFHub. It's only when I learn TF tutorials like 4 yeras ago, and it's still TF 1 age... As far as I remember, it's convenient to use as a text embedding layer, but not as easy changing the loaded model detail. Just checked the hub, and it has been enriched a lot since then. May check some of those models if got time.