tf-extra 1.1.1

A powerful extension library for TensorFlow/Keras — MultiOptimizer, custom layers, losses, and callbacks.

🚀 TF-Extra

A powerful extension library for TensorFlow/Keras — the spiritual successor to tensorflow_addons.

---

✨ Features

Module	Components
🔧 Optimizers	MultiOptimizer — assign different optimizers to different layers
📐 Layers	CBAM, SqueezeExcitation, MultiHeadSelfAttention, GeM, Mish, PositionalEncoding, AdaptiveConcatPool2D
📉 Losses	FocalLoss, DiceLoss, TverskyLoss, LabelSmoothingCrossEntropy, ContrastiveLoss, TripletLoss
📞 Callbacks	CosineAnnealingScheduler, GradientAccumulationCallback, TimeLimitCallback

---

📦 Installation

pip install tf-extra

Or install from source:

git clone https://github.com/yourusername/tf-extra.git
cd tf-extra
pip install -e .

---

🚀 Quick Start

MultiOptimizer (Drop-in replacement for tfa.optimizers.MultiOptimizer)

import tensorflow as tf
import tf_extra

# Load a pre-trained model
base_model = tf.keras.applications.ResNet50(weights='imagenet', include_top=False)
x = tf.keras.layers.GlobalAveragePooling2D()(base_model.output)
x = tf.keras.layers.Dense(1)(x)
model = tf.keras.Model(base_model.input, x)

# Define different learning rates for different parts of the model
optimizers = [
    tf.keras.optimizers.Adam(learning_rate=1e-6),   # Very slow for early layers
    tf.keras.optimizers.Adam(learning_rate=1e-5),   # Medium for middle layers
    tf.keras.optimizers.Adam(learning_rate=1e-4),   # Fast for the head
]

optimizers_and_layers = [
    (optimizers[0], base_model.layers[:100]),
    (optimizers[1], base_model.layers[100:]),
    (optimizers[2], model.layers[-1]),
]

opt = tf_extra.optimizers.MultiOptimizer(optimizers_and_layers)
model.compile(
    optimizer=opt,
    loss=tf.keras.losses.BinaryCrossentropy(from_logits=True),
    metrics=[tf.keras.metrics.BinaryAccuracy(threshold=0, name='accuracy')]
)

Focal Loss (for imbalanced datasets)

import tf_extra

model.compile(
    optimizer='adam',
    loss=tf_extra.losses.FocalLoss(alpha=0.25, gamma=2.0),
    metrics=['accuracy']
)

CBAM Attention Layer

import tf_extra

x = tf.keras.layers.Conv2D(64, 3, padding='same')(inputs)
x = tf_extra.layers.CBAM(reduction_ratio=16)(x)

Cosine Annealing with Warm Restarts

import tf_extra

scheduler = tf_extra.callbacks.CosineAnnealingScheduler(
    eta_min=1e-7,
    eta_max=1e-3,
    T_0=10,
    T_mult=2
)
model.fit(X, y, callbacks=[scheduler])

---

📖 Full API Reference

tf_extra.optimizers.MultiOptimizer

MultiOptimizer(
    optimizers_and_layers,     # List of (optimizer, layers) tuples
    fallback_optimizer=None,   # Optional: optimizer for unassigned variables
    name="MultiOptimizer"
)

tf_extra.layers

Layer	Description
CBAM(reduction_ratio=16, kernel_size=7)	Convolutional Block Attention Module
SqueezeExcitation(reduction_ratio=16)	SE-Net attention block
MultiHeadSelfAttention(embed_dim, num_heads)	Transformer-style self-attention
GeM(p_init=3.0)	Generalized Mean Pooling
Mish()	Mish activation function
PositionalEncoding(max_len=5000)	Sinusoidal positional encoding
AdaptiveConcatPool2D()	Concatenated average + max pooling
SpatialDropout1D(rate=0.5)	Spatial dropout for 1D data

tf_extra.losses

Loss	Description
FocalLoss(alpha=0.25, gamma=2.0)	Focal loss for class imbalance
DiceLoss(smooth=1.0)	Dice loss for segmentation
TverskyLoss(alpha=0.5, beta=0.5)	Generalized Dice loss
LabelSmoothingCrossEntropy(smoothing=0.1)	Label smoothing
ContrastiveLoss(margin=1.0)	For Siamese networks
TripletLoss(margin=1.0)	For metric learning

tf_extra.callbacks

Callback	Description
CosineAnnealingScheduler(eta_min, eta_max, T_0, T_mult)	SGDR scheduler
GradientAccumulationCallback(accumulation_steps=4)	Gradient accumulation
TimeLimitCallback(max_seconds)	Stop training after time limit

---

🧪 Testing

pip install -e ".[dev]"
pytest tests/ -v

---

🤝 Contributing

Contributions are welcome! Please read the contributing guidelines and submit a PR.

---

📄 License

MIT License — see LICENSE for details.

---

🙏 Acknowledgements

Inspired by tensorflow_addons which is no longer maintained. Built to provide the same functionality with modern TensorFlow compatibility.