jina-ai / finetuner

Finetuning any deep neural network for better embedding on neural search tasks

Finetuner allows one to tune the weights of any deep neural network for better embeddings on search tasks. It accompanies Jina to deliver the last mile of performance for domain-specific neural search applications.

🎛 Designed for finetuning: a human-in-the-loop deep learning tool for leveling up your pretrained models in domain-specific neural search applications.

🔱 Powerful yet intuitive: all you need is finetuner.fit() - a one-liner that unlocks rich features such as siamese/triplet network, metric learning, self-supervised pretraining, layer pruning, weights freezing, dimensionality reduction.

⚛️ Framework-agnostic: promise an identical API & user experience on PyTorch, Tensorflow/Keras and PaddlePaddle deep learning backends.

🧈 DocArray integration: buttery smooth integration with DocArray, reducing the cost of context-switch between experiment and production.

How does it work

Install

Requires Python 3.7+ and one of PyTorch(>=1.9) or Tensorflow(>=2.5) or PaddlePaddle installed on Linux/MacOS.

pip install finetuner

Finetuning ResNet50 on CelebA

1. Download CelebA-small dataset (7.7MB) and decompress it to './img_align_celeba'. Full dataset can be found here.
2. Finetuner accepts docarray DocumentArray, so we load CelebA image into this format using a generator:

   from docarray import DocumentArray
   
   # please change the file path to your data path
   data = DocumentArray.from_files('img_align_celeba/*.jpg')
   
   
   def preproc(doc):
       return (
           doc.load_uri_to_image_tensor(224, 224)
           .set_image_tensor_normalization()
           .set_image_tensor_channel_axis(-1, 0)
       )  # No need for normalization and changing channel axes line if you are using tf/keras
   
   
   data.apply(preproc)

3. Load pretrained ResNet18 using PyTorch/Keras/Paddle:
   • PyTorch

     import torchvision
     resnet = torchvision.models.resnet50(pretrained=True)

   • Keras

     import tensorflow as tf
     resnet = tf.keras.applications.resnet50.ResNet50(weights='imagenet')

   • Paddle

     import paddle
     resnet = paddle.vision.models.resnet50(pretrained=True)

4. Start the Finetuner:

   import finetuner as ft
   
   tuned_model = ft.fit(
       model=resnet,
       train_data=data,
       loss='TripletLoss',
       epochs=20,
       device='cuda',
       batch_size=128,
       to_embedding_model=True,
       input_size=(3, 224, 224), # for keras use (224, 224, 3)
       freeze=False,
   )

Support

• Check out the Learning Bootcamp to get started with Finetuner.
• Use Discussions to talk about your use cases, questions, and support queries.
• Join our Slack community and chat with other Jina community members about ideas.
• Join our Engineering All Hands meet-up to discuss your use case and learn Jina's new features.
  • When? The second Tuesday of every month
  • Where? Zoom (see our public events calendar/.ical) and live stream on YouTube
• Subscribe to the latest video tutorials on our YouTube channel

Join Us

Finetuner is backed by Jina AI and licensed under Apache-2.0. We are actively hiring AI engineers, solution engineers to build the next neural search ecosystem in opensource.