Quantized ImageBind

[ahmedsaed]

This pull request introduces quantization support for the ImageBind model, enabling efficient inference with reduced precision. The changes include adding new quantization-related classes, modifying existing components to support quantization, and providing a script for model quantization and evaluation.

Quantization Support Enhancements:

• New Quantization Classes:

  • Added QuantizableMultiheadAttention and QuantizedMultiheadAttention classes to support quantized attention mechanisms in the transformer module (imagebind/models/transformer.py). [1] [2]
  • Introduced QuantizedDropPath to handle quantization-aware stochastic depth. This class includes quantization and dequantization stubs for compatibility with quantized models (imagebind/models/transformer.py).

• Integration of Quantization in Model Components:

  • Replaced DropPath with QuantizedDropPath and added nn.quantized.FloatFunctional for skip connections in the transformer blocks (imagebind/models/transformer.py). [1] [2]
  • Substituted MultiheadAttention with QuantizableMultiheadAttention in the model's trunk instantiation (imagebind/models/imagebind_model.py).

Codebase Refactoring for Readability:

• Reformatted import statements for better readability and consistency across files (imagebind/models/imagebind_model.py).
• Simplified string formatting in extra_repr for the LearnableLogitScaling class (imagebind/models/helpers.py).

Quantization Workflow Script:

• Added a new script (quantized.py) to demonstrate the quantization process, including:
  • Preparing the ImageBind model for static quantization.
  • Creating dummy data for calibration.
  • Evaluating the similarity between original and quantized model outputs using cosine similarity.

Minor Improvements:

• Added a check to dequantize inputs if they are quantized before normalization in the Normalize class (imagebind/models/helpers.py).
• Adjusted model weight loading to improve code readability (imagebind/models/imagebind_model.py).