Introduce kimm: a collection of image models, blocks and layers written in Keras3

[james77777778]

Repository: https://github.com/james77777778/keras-image-models

If you want to try a different model that isn't available in Keras, you can explore this library.
Feel free to contact me if you encounter any issues or have ideas.

Introduction

Keras Image Models (kimm) aims to provide SOTA models with pretrained weights in a user-friendly manner.

Features:

• 🚀 A model zoo where almost all models come with pre-trained weights on ImageNet.
• 🧰 Providing APIs to export models to .tflite and .onnx.
• 🔧 Supporting the reparameterization technique.
• ✨ Integrated with feature extraction capability.

Usage

import keras
import kimm
import numpy as np


# List available models
print(kimm.list_models("mobileone", weights="imagenet"))  # ['MobileOneS0', ...]

# Initialize model with pretrained ImageNet weights
x = keras.random.uniform([1, 224, 224, 3])
model = kimm.models.MobileOneS0()
y = model.predict(x)
print(y.shape)  # (1, 1000)

# Get reparameterized model by kimm.utils.get_reparameterized_model
reparameterized_model = kimm.utils.get_reparameterized_model(model)
y2 = reparameterized_model.predict(x)
np.testing.assert_allclose(
    keras.ops.convert_to_numpy(y), keras.ops.convert_to_numpy(y2), atol=1e-5
)

# Export model to tflite format
kimm.export.export_tflite(reparameterized_model, 224, "model.tflite")

# Export model to onnx format (note: must be "channels_first" format)
# kimm.export.export_onnx(reparameterized_model, 224, "model.onnx")

# Enable feature extraction by setting `feature_extractor=True`
# `feature_keys` can be optionally specified
model = kimm.models.MobileOneS0(
    include_top=False, feature_extractor=True, feature_keys=["BLOCK2_S16", "BLOCK3_S32"]
)
features = model.predict(x)
for feature_name, feature in features.items():
    print(feature_name, feature.shape)  # BLOCK2_S16 (1, 14, 14, 256) ...

Quickstart

An end-to-end example: fine-tuning an image classification model on a cats vs. dogs dataset

Using kimm.models.EfficientNetLiteB0:

Reference: Transfer learning & fine-tuning (keras.io)

Grad-CAM

Using kimm.models.MobileViTS:

Reference: Grad-CAM class activation visualization (keras.io)

Model Zoo

Model	Paper	Weights are ported from	API
ConvMixer	ICLR 2022 Submission	timm	kimm.models.ConvMixer*
ConvNeXt	CVPR 2022	timm	kimm.models.ConvNeXt*
DenseNet	CVPR 2017	timm	kimm.models.DenseNet*
EfficientNet	ICML 2019	timm	kimm.models.EfficientNet*
EfficientNetLite	ICML 2019	timm	kimm.models.EfficientNetLite*
EfficientNetV2	ICML 2021	timm	kimm.models.EfficientNetV2*
GhostNet	CVPR 2020	timm	kimm.models.GhostNet*
GhostNetV2	NeurIPS 2022	timm	kimm.models.GhostNetV2*
HGNet		timm	kimm.models.HGNet*
HGNetV2		timm	kimm.models.HGNetV2*
InceptionNeXt	arXiv 2023	timm	kimm.models.InceptionNeXt*
InceptionV3	CVPR 2016	timm	kimm.models.InceptionV3
LCNet	arXiv 2021	timm	kimm.models.LCNet*
MobileNetV2	CVPR 2018	timm	kimm.models.MobileNetV2*
MobileNetV3	ICCV 2019	timm	kimm.models.MobileNetV3*
MobileOne	CVPR 2023	timm	kimm.models.MobileOne*
MobileViT	ICLR 2022	timm	kimm.models.MobileViT*
MobileViTV2	arXiv 2022	timm	kimm.models.MobileViTV2*
RegNet	CVPR 2020	timm	kimm.models.RegNet*
RepVGG	CVPR 2021	timm	kimm.models.RepVGG*
ResNet	CVPR 2015	timm	kimm.models.ResNet*
TinyNet	NeurIPS 2020	timm	kimm.models.TinyNet*
VGG	ICLR 2015	timm	kimm.models.VGG*
ViT	ICLR 2021	timm	kimm.models.VisionTransformer*
Xception	CVPR 2017	keras	kimm.models.Xception

[ibinti]

Thanks for your great work.
I am very eager to try out MobileViTV2 with ported timm imagenet weights,
Following MobileOneS0 example, I did

kimm.list_models("mobilevit2", weights="imagenet") 
=> ['MobileViTV2W050', 'MobileViTV2W075', 'MobileViTV2W100', 'MobileViTV2W125', 'MobileViTV2W150', 'MobileViTV2W175', 'MobileViTV2W200']

x = keras.random.uniform([1, 224, 224, 3])
model = kimm.models.MobileViTV2W100()
y = model.predict(x)
print(y.shape)

However, model.predict(x) throws

TypeError: Exception encountered when calling Reshape.call().
cannot reshape array of shape (1, 28, 28, 128) (size 100352) into shape [-1, 16, 2, 16, 2, 128] because the product of specified axis sizes (131072) does not evenly divide 100352
Arguments received by Reshape.call():
  • x=jnp.ndarray(shape=(1, 28, 28, 128), dtype=float32)

Thanks

[ibinti]

Well, I found the reason. MobileViTV2 expects image shape (255,255,3)
x = keras.random.uniform([1, 256, 256, 3])
Thanks

[ibinti]

I wonder how to handle image resizing and scaling. Is there a way to retrieve what the model expects like shape, pixel value range [0,255] float or int , or [0,1] or [-1,1]?

[james77777778]

Hey @ibinti
Thanks for trying out this project.

Actually, there was an update a few days ago to support arbitrary sizes for MobileViT and other ViT models.
https://github.com/james77777778/keras-image-models/releases/tag/0.2.1

I wonder how to handle image resizing and scaling. Is there a way to retrieve what the model expects like shape, pixel value range [0,255] float or int , or [0,1] or [-1,1]?

You can get the default size by:

model = kimm.models.MobileNetV2W050()
print(model._default_size)  # prints 224

You can specify an arbitrary size for the model:

# `input_shape` manner
x = keras.random.uniform([1, 256, 256, 3])
model = kimm.models.MobileNetV2W050(input_shape=[256, 256, 3])
y = model.predict(x)

# `input_tensor` manner
x = keras.random.uniform([1, 256, 256, 3])
inputs = keras.Input([256, 256, 3])
model = kimm.models.MobileNetV2W050(input_tensor=inputs)
y = model.predict(x)

All kimm models expect the value range of [0, 255]. The details of preprocessing can only be inspected in the source code.

• https://github.com/james77777778/keras-image-models/blob/main/kimm/_src/models/mobilevit.py#L571
• https://github.com/james77777778/keras-image-models/blob/main/kimm/_src/models/base_model.py#L147

Maybe I can add this information to the models in the future. Thanks for pointint that out!

[ibinti]

Great. Now I am ready to train kimm.models.MobileViTV2W100 and compare its performance with keras.applications.EfficientNetV2.
Thanks again!

[BWBrook]

Hi James. I've been using pretrained kimm models successfully in Docker containers. After training a model via fine-tuning and saving it to disk as a .keras file in one Docker container, if I load it up into a different Docker container to be used for inference, it triggers an automatic download of the pretrained weights again, even though they are not required for inference and the model frozen prior to saving it.

Is there a way to avoid this behaviour? This doesn't happen if I use the keras.applications API for the same pretrained model. But I want to use kimm because it has a larger variety of available models!

[BWBrook]

Thanks for answering james. Yes, I think kimm is great. But I didn't mean loading weights when building a model. For instance, when I'm training a new model I can use something like:

model = kimm.models.EfficientNetV2S(include_top=False)
This gives me pretrained imagenet weights (which I want) and I can then fine-tune on this, and then save the fine tuned model it with:
saving.save_model(model, model_fine_tuned.keras, include_optimizer=False)

That all works fine and saves a .keras model with the fine-tuned weights included.

But if I then try to load model_fine_tuned.keras in a new program (e.g., new Docker image) using this:
model = saving.load_model("/code/model_fine_tuned.keras", compile=False)
Then the model loads fine, and can be used for inference as expected BUT the load_model call triggers a download of the imagenet weights again, even though it doesn't need to. If I'd used a keras.applications.EfficientNetV2S instead, this download does not happen.

So I was wondering how to avoid kimm models triggering the re-downloading of the weights when they are loaded using load_model. Do you see what I'm saying? :)

[james77777778]

Oh I see. I have verified that there is a redownloading issue caused by the config json.
A simple fix is on the way and I will push a new version soon.

james77777778/keras-image-models#54

[james77777778]

@BWBrook
I believe kimm==2.5.0 should fix the problem. It should be available on pypi now.

[BWBrook]

Thanks @james77777778 that's excellent. I can confirm that this has indeed fixed the re-loading problem! I really appreciate your prompt intervention. In case you're interested, this is where I'm deploying kimm: MEWC - Mega Efficient Wildlife Classifier

[james77777778]

Wow! It’s great to hear that kimm can be integrated into such an amazing project.

Thank you for sharing.