Search code, repositories, users, issues, pull requests...

• Introduction • Installation • Usage • Examples • Available Architectures • Contributing • Acknowledgements

flaim is a library of state-of-the-art pre-trained vision models, plus common deep learning modules in computer vision, for Flax. It exposes a host of diverse image models through a straightforward interface with an emphasis on simplicity, leanness, and readability, and supplies lower-level modules for designing custom architectures.

flaim can be installed through pip install flaim. Beware that pip installs the CPU version of JAX, and you must manually install JAX yourself to run your programs on a GPU or TPU.

pretrained (Union[str, int, bool]): Every flaim network is accompanied by at least one group of pre-trained parameters. For example, those of MaxViT-Small (maxvit_small) are in1k_224, in1k_384, and in1k_512, corresponding to parameters trained on ImageNet1K at resolutions 224 x 224, 384 x 384, and 512 x 512 respectively. When pretrained is

input_size (int): When pretrained is False, input_size refers to the input size the model should expect and is used to initialize the parameters. Providing the correct value for input_size is especially important for fixed-resolution architectures such as ViT.

Leave a Comment