<jupyter_start><jupyter_code>!pip install --upgrade tensorflow
!pip install "git+https://github.com/DavidLandup0/keras-cv@vit"
import math
import sys
import os
import pandas as pd

import tensorflow as tf
import keras_cv

import matplotlib.pyplot as plt
import numpy as np
config_names = {
    # JAX name : KCV name, layer_nums
    "Ti/16": ("ViTTiny16", 12),
    "S/16": ("ViTS16", 12),
    "B/16": ("ViTB16", 12),
    "L/16": ("ViTL16", 24),
    "S/32": ("ViTS32", 12),
    "B/32": ("ViTB32", 12),
}
# Choose model to convert
model_to_convert = list(config_names.items())[0]
model_to_convert
model = eval(
    f"keras_cv.models.{model_to_convert[1][0]}(include_rescaling=False, include_top=True, num_classes=1000, weights=None, input_shape=(224, 224, 3))"
)
with tf.io.gfile.GFile("gs://vit_models/augreg/index.csv") as f:
    df = pd.read_csv(f)

df.head()
model_df = df.query(
    f'ds=="i21k" & adapt_resolution==224 & adapt_ds=="imagenet2012" & name=="{model_to_convert[0]}"'
).sort_values("adapt_final_test", ascending=False)
model_df.head()
best_model_i1k_checkpoint = str(model_df.iloc[0]["adapt_filename"])
model_df.iloc[0]["adapt_filename"], model_df.iloc[0]["adapt_final_test"]
filename = best_model_i1k_checkpoint

path = f"gs://vit_models/augreg/{filename}.npz"

print(f"{tf.io.gfile.stat(path).length / 1024 / 1024:.1f} MiB - {path}")
local_path = path.split("//")[-1].split("/")[-1]
local_path
!gsutil cp {path} .
with open(local_path, "rb") as f:
    params_jax = np.load(f)
    params_jax = dict(zip(params_jax.keys(), params_jax.values()))
from pprint import pformat

print(pformat(list(params_jax.keys())))
jax_params_to_kcv_params = {
    "Transformer/posembed_input/pos_embedding": "patch_embedding/embedding/embeddings",
    "embedding/bias": "patch_embedding_1/dense_26/bias",
    "embedding/kernel": "patch_embedding_1/dense_26/kernel",
    "cls": "patch_embedding_1/class_token",
    "Transformer/encoderblock_0/LayerNorm_0/scale": "transformer_encoder_12/layer_normalization_25/gamma",
    "Transformer/encoderblock_0/LayerNorm_0/bias": "transformer_encoder_12/layer_normalization_25/beta",
    "Transformer/encoderblock_0/LayerNorm_2/scale": "transformer_encoder_12/layer_normalization_26/gamma",
    "Transformer/encoderblock_0/LayerNorm_2/bias": "transformer_encoder_12/layer_normalization_26/beta",
    "Transformer/encoderblock_0/MultiHeadDotProductAttention_1/query/kernel": "transformer_encoder_12/multi_head_attention_12/query/kernel",
    "Transformer/encoderblock_0/MultiHeadDotProductAttention_1/query/bias": "transformer_encoder_12/multi_head_attention_12/query/bias",
    "Transformer/encoderblock_0/MultiHeadDotProductAttention_1/key/kernel": "transformer_encoder_12/multi_head_attention_12/key/kernel",
    "Transformer/encoderblock_0/MultiHeadDotProductAttention_1/key/bias": "transformer_encoder_12/multi_head_attention_12/key/bias",
    "Transformer/encoderblock_0/MultiHeadDotProductAttention_1/value/kernel": "transformer_encoder_12/multi_head_attention_12/value/kernel",
    "Transformer/encoderblock_0/MultiHeadDotProductAttention_1/value/bias": "transformer_encoder_12/multi_head_attention_12/value/bias",
    "Transformer/encoderblock_0/MultiHeadDotProductAttention_1/out/kernel": "transformer_encoder_12/multi_head_attention_12/attention_output/kernel",
    "Transformer/encoderblock_0/MultiHeadDotProductAttention_1/out/bias": "transformer_encoder_12/multi_head_attention_12/attention_output/bias",
    "Transformer/encoderblock_0/MlpBlock_3/Dense_0/kernel": "transformer_encoder_12/dense_27/kernel",
    "Transformer/encoderblock_0/MlpBlock_3/Dense_0/bias": "transformer_encoder_12/dense_27/bias",
    "Transformer/encoderblock_0/MlpBlock_3/Dense_1/kernel": "transformer_encoder_12/dense_28/kernel",
    "Transformer/encoderblock_0/MlpBlock_3/Dense_1/bias": "transformer_encoder_12/dense_28/bias",
    # ... other transformer blocks
    "Transformer/encoder_norm/scale": "layer_normalization_49/gamma",
    "Transformer/encoder_norm/bias": "layer_normalization_49/beta",
}
model.layers
model.summary()
# Check shapes for the class token and embedding layers
print(params_jax["cls"].shape)
print(params_jax["embedding/kernel"].shape)
print(params_jax["embedding/bias"].shape)
print(params_jax["Transformer/posembed_input/pos_embedding"].shape)

for w in model.layers[1].weights:
    print(w.name, w.shape)
# Copy PatchingAndEmbedding layer
model.layers[1].weights[0].assign(tf.Variable(params_jax["cls"]))
model.layers[1].weights[1].assign(tf.Variable(params_jax["embedding/kernel"]))
model.layers[1].weights[2].assign(tf.Variable(params_jax["embedding/bias"]))
model.layers[1].weights[3].assign(
    tf.Variable(
        params_jax["Transformer/posembed_input/pos_embedding"].squeeze()
    )
)
# Check transformer block shapes between JAX and KCV
print(params_jax["Transformer/encoderblock_4/LayerNorm_0/scale"].shape)
print(params_jax["Transformer/encoderblock_4/LayerNorm_0/bias"].shape)
print(params_jax["Transformer/encoderblock_4/LayerNorm_2/scale"].shape)
print(params_jax[f"Transformer/encoderblock_4/LayerNorm_2/bias"].shape)
print(
    params_jax[
        f"Transformer/encoderblock_4/MultiHeadDotProductAttention_1/query/kernel"
    ].shape
)
print(
    params_jax[
        f"Transformer/encoderblock_4/MultiHeadDotProductAttention_1/query/bias"
    ].shape
)
print(
    params_jax[
        f"Transformer/encoderblock_4/MultiHeadDotProductAttention_1/key/kernel"
    ].shape
)
print(
    params_jax[
        f"Transformer/encoderblock_4/MultiHeadDotProductAttention_1/key/bias"
    ].shape
)
print(
    params_jax[
        f"Transformer/encoderblock_4/MultiHeadDotProductAttention_1/value/kernel"
    ].shape
)
print(
    params_jax[
        f"Transformer/encoderblock_4/MultiHeadDotProductAttention_1/value/bias"
    ].shape
)
print(
    params_jax[
        f"Transformer/encoderblock_4/MultiHeadDotProductAttention_1/out/kernel"
    ].shape
)
print(
    params_jax[
        f"Transformer/encoderblock_4/MultiHeadDotProductAttention_1/out/bias"
    ].shape
)
print(params_jax[f"Transformer/encoderblock_4/MlpBlock_3/Dense_0/kernel"].shape)
print(params_jax[f"Transformer/encoderblock_4/MlpBlock_3/Dense_0/bias"].shape)
print(params_jax[f"Transformer/encoderblock_4/MlpBlock_3/Dense_1/kernel"].shape)
print(params_jax[f"Transformer/encoderblock_4/MlpBlock_3/Dense_1/bias"].shape)

for w in model.layers[4].weights:
    print(w.name, w.shape)
# Copy Transformer Encoders
for i in range(model_to_convert[1][1]):
    model.layers[3 + i].weights[0].assign(
        tf.Variable(
            params_jax[f"Transformer/encoderblock_{i}/LayerNorm_0/scale"]
        )
    )
    model.layers[3 + i].weights[1].assign(
        tf.Variable(
            params_jax[f"Transformer/encoderblock_{i}/LayerNorm_0/bias"]
        )
    )
    model.layers[3 + i].weights[2].assign(
        tf.Variable(
            params_jax[f"Transformer/encoderblock_{i}/LayerNorm_2/scale"]
        )
    )
    model.layers[3 + i].weights[3].assign(
        tf.Variable(
            params_jax[f"Transformer/encoderblock_{i}/LayerNorm_2/bias"]
        )
    )
    model.layers[3 + i].weights[4].assign(
        tf.Variable(
            params_jax[
                f"Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/query/kernel"
            ]
        )
    )
    model.layers[3 + i].weights[5].assign(
        tf.Variable(
            params_jax[
                f"Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/query/bias"
            ]
        )
    )
    model.layers[3 + i].weights[6].assign(
        tf.Variable(
            params_jax[
                f"Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/key/kernel"
            ]
        )
    )
    model.layers[3 + i].weights[7].assign(
        tf.Variable(
            params_jax[
                f"Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/key/bias"
            ]
        )
    )
    model.layers[3 + i].weights[8].assign(
        tf.Variable(
            params_jax[
                f"Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/value/kernel"
            ]
        )
    )
    model.layers[3 + i].weights[9].assign(
        tf.Variable(
            params_jax[
                f"Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/value/bias"
            ]
        )
    )
    model.layers[3 + i].weights[10].assign(
        tf.Variable(
            params_jax[
                f"Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/out/kernel"
            ]
        )
    )
    model.layers[3 + i].weights[11].assign(
        tf.Variable(
            params_jax[
                f"Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/out/bias"
            ].reshape(model.layers[3 + i].weights[11].shape)
        )
    )
    model.layers[3 + i].weights[12].assign(
        tf.Variable(
            params_jax[
                f"Transformer/encoderblock_{i}/MlpBlock_3/Dense_0/kernel"
            ]
        )
    )
    model.layers[3 + i].weights[13].assign(
        tf.Variable(
            params_jax[f"Transformer/encoderblock_{i}/MlpBlock_3/Dense_0/bias"]
        )
    )
    model.layers[3 + i].weights[14].assign(
        tf.Variable(
            params_jax[
                f"Transformer/encoderblock_{i}/MlpBlock_3/Dense_1/kernel"
            ]
        )
    )
    model.layers[3 + i].weights[15].assign(
        tf.Variable(
            params_jax[f"Transformer/encoderblock_{i}/MlpBlock_3/Dense_1/bias"]
        )
    )
print(params_jax["Transformer/encoder_norm/scale"].shape)
print(params_jax["Transformer/encoder_norm/bias"].shape)

for w in model.layers[15].weights:
    print(w.name, w.shape)
# Copy layer norm before class head
model.layers[15].weights[0].assign(
    tf.Variable(params_jax["Transformer/encoder_norm/scale"])
)
model.layers[15].weights[1].assign(
    tf.Variable(params_jax["Transformer/encoder_norm/bias"])
)
print(params_jax["head/kernel"].shape)
print(params_jax["head/bias"].shape)

for w in model.layers[17].weights:
    print(w.name, w.shape)
# Copy haed kernel and bias
model.layers[17].weights[0].assign(tf.Variable(params_jax["head/kernel"]))
model.layers[17].weights[1].assign(tf.Variable(params_jax["head/bias"]))
import matplotlib.pyplot as plt
import cv2
import numpy as np

import PIL
import urllib


def url_to_array(url):
    req = urllib.request.urlopen(url)
    arr = np.array(bytearray(req.read()), dtype=np.int8)
    arr = cv2.imdecode(arr, -1)
    arr = cv2.cvtColor(arr, cv2.COLOR_BGR2RGB)
    arr = cv2.resize(arr, (224, 224))
    return arr


def preprocess_image(image, label):
    image_resized = tf.image.resize(image, (224, 224))
    image_resized = tf.cast(image_resized, tf.float32)
    image_resized = (image_resized - 127.5) / 127.5
    return image_resized, label
cat = "https://upload.wikimedia.org/wikipedia/commons/thumb/4/4d/Cat_November_2010-1a.jpg/1200px-Cat_November_2010-1a.jpg"
cat_img = url_to_array(cat)
cat_img, _ = preprocess_image(cat_img, None)
cat_img = tf.expand_dims(cat_img, 0)
!wget https://storage.googleapis.com/bit_models/ilsvrc2012_wordnet_lemmas.txt -O ilsvrc2012_wordnet_lemmas.txt

with open("ilsvrc2012_wordnet_lemmas.txt", "r") as f:
    lines = f.readlines()
imagenet_int_to_str = [line.rstrip() for line in lines]
predictions = model.predict(cat_img)
top_5 = tf.math.top_k(predictions, k=5, sorted=False)
top_5
pred = np.argmax(predictions)
imagenet_int_to_str[int(pred)]
dog_url = "https://hips.hearstapps.com/hmg-prod.s3.amazonaws.com/images/dog-puppy-on-garden-royalty-free-image-1586966191.jpg?crop=1.00xw:0.669xh;0,0.190xh&resize=640:*"
dog_img = url_to_array(dog_url)
dog_img, _ = preprocess_image(dog_img, None)
dog_img = tf.expand_dims(dog_img, 0)
predictions = model.predict(dog_img)
pred = np.argmax(predictions)
imagenet_int_to_str[int(pred)]
model.compile(
    "adam",
    "sparse_categorical_crossentropy",
    metrics=["accuracy", keras.metrics.SparseTopKCategoricalAccuracy(5)],
)
import tensorflow_datasets as tfds

(test_set), info = tfds.load(
    "imagenet_v2", split=["test"], as_supervised=True, with_info=True
)
test_set = (
    test_set[0]
    .shuffle(len(test_set[0]))
    .map(preprocess_image)
    .batch(32)
    .prefetch(tf.data.AUTOTUNE)
)
for entry, label in test_set.take(1):
    print(label)
model.evaluate(test_set)
model.save(f"{model_to_convert[1][0]}.h5")<jupyter_output><empty_output>