Image embedder finetuning example

Hello i try to reproduce the finetuning image embedder example https://lightning-flash.readthedocs.io/en/latest/reference/image_embedder.html#finetuning

import flash
from flash import download_data
from flash.vision import ImageClassificationData, ImageEmbedder

# 1. Download the data
download_data("https://pl-flash-data.s3.amazonaws.com/hymenoptera_data.zip", "data/")

# 2. Load the data
datamodule = ImageClassificationData.from_folders(
    train_folder="data/hymenoptera_data/train/",
    valid_folder="data/hymenoptera_data/val/",
    test_folder="data/hymenoptera_data/test/",
)

# 3. Build the model
embedder = ImageEmbedder(backbone="resnet18", embedding_dim=128)

# 4. Create the trainer. Run once on data
trainer = flash.Trainer(max_epochs=1)

# 5. Train the model
trainer.finetune(embedder, datamodule=datamodule, strategy="freeze_unfreeze")

# 6. Test the model
trainer.test()

# 7. Save it!
trainer.save_checkpoint("image_embedder_model.pt")

After that i use the saved checkpoint :

model = ImageEmbedder.load_from_checkpoint("image_embedder_model.pt")
model.predict("data/hymenoptera_data/predict/")

But the result is a class predictions instead of a list of vectors : [[[0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]]

lightning-flash 0.2.1

[psardin]

I have changed the backbone to swav-imagenet and the embeddings are much better for my use case. Im closing the issue.

[edgarriba]

@psardin image embeddings can be trained in many ways. In this example it’s shown how to train them based on an image classification loss - that’s the reason of using ImageClassificationData and why the return value as the one hot encoding of the predicted classes. As @aniketmaurya suggested, if you could be more specific about your application we could suggest an alternative.