pip install -U keras-visionpip install git+https://github.com/veb-101/keras-vision.git1import os
2
3os.environ["KERAS_BACKEND"] = "jax"
4# os.environ["KERAS_BACKEND"] = "tensorflow"
5# os.environ["KERAS_BACKEND"] = "torch"
6
7from keras_vision.fastvit import fastvit_ma36
8
9model = fastvit_ma36(
10 pretrained=True,
11 inference_mode=True,
12 load_kd_weights=True,
13)
14model.summary()keras_fastvit_ma36.weights.h5 - Weights for model variant - ma36keras_fastvit_ma36_reparam.weights.h5 - Reparameterized weights for model variant - ma36 (inference_mode = True)keras_fastvit_ma36_kd.weights.h5 - Weights for model variant - ma36, trained using knowledge distillation.keras_fastvit_ma36_reparam_kd.weights.h5 - Reparameterized weights for model variant - ma36, trained using knowledge distillation (inference_mode = True)