- Implementation of the basic Vision Transformer using
jax
anddm-haiku
- Model and optimizer serialization/deserialization with pickle
- Uses tensor manipulation operations using
einops
- Re-implementation of Attention Rollout with percentile based filtering from
https://github.com/jacobgil/vit-explain
(WIP) - Inference script with attention rollout (WIP)
dogeplusplus / haiku-transformer Goto Github PK
View Code? Open in Web Editor NEWVision Transformer implemented with JAX and dm-haiku
License: MIT License