PyTorch and Zilliz Cloud Integration

PyTorch est une bibliothèque d’apprentissage automatique open source développée par le laboratoire Meta AI Research et faisant désormais partie de la Linux Foundation. Elle offre des capacités de calcul flexibles pour le deep learning avec des API intuitives, une accélération GPU et un écosystème complet. La bibliothèque prend en charge divers domaines de l’apprentissage automatique, notamment la vision par ordinateur, le NLP et l’apprentissage par renforcement, avec comme fonctionnalité clé ses graphes de calcul dynamiques pouvant être modifiés pendant l’exécution, permettant une expérimentation et un prototypage rapides.

En s’intégrant à Zilliz Cloud (Milvus entièrement géré), les embeddings générés par PyTorch à partir de plusieurs types de données — images, texte, audio — peuvent être stockés, gérés et interrogés via une base de données vectorielle évolutive et efficace, permettant des applications telles que les moteurs de recherche d’images, les systèmes de recommandation et les outils de détection d’anomalies grâce à une recherche de similarité vectorielle haute performance.

• Génération d’embeddings flexible avec stockage évolutif : Les modèles préentraînés de PyTorch (p. ex., ResNet50) génèrent des vecteurs de caractéristiques de haute qualité à partir de divers types de données, tandis que Zilliz Cloud stocke et indexe ces embeddings pour une recherche de similarité rapide à grande échelle.
• Pipeline accéléré par GPU : PyTorch exploite l’accélération GPU pour une génération efficace d’embeddings, et Zilliz Cloud fournit une récupération vectorielle à faible latence, créant ainsi un pipeline de bout en bout hautes performances.
• Écosystème de modèles riche : La bibliothèque Torchvision de PyTorch fournit des modèles préentraînés pour les tâches de vision par ordinateur, ce qui facilite la création d’applications de recherche d’images et de similarité visuelle adossées au stockage vectoriel de Zilliz Cloud.
• Prise en charge du traitement par lots : L’intégration prend en charge la génération et l’insertion d’embeddings par lots, permettant un traitement efficace de jeux de données à grande échelle dans Zilliz Cloud.

PyTorch sert de framework d’apprentissage profond, générant des embeddings de vecteurs de caractéristiques à partir de données brutes à l’aide de modèles pré-entraînés. Pour la recherche d’images, il utilise des modèles comme ResNet50 de la bibliothèque Torchvision afin d’extraire des représentations de caractéristiques à haute dimension à partir d’images, avec la couche de classification finale supprimée pour produire de purs vecteurs d’embedding.

Zilliz Cloud sert de couche de base de données vectorielle, stockant et indexant les embeddings générés par PyTorch. Il fournit une recherche de similarité haute performance à l’aide de métriques comme la distance L2, permettant une récupération rapide des éléments les plus similaires à partir de grandes collections.

Ensemble, PyTorch et Zilliz Cloud créent une solution complète de recherche de similarité : PyTorch traite les données brutes (images, texte, audio) en embeddings vectoriels à l’aide de modèles pré-entraînés, et Zilliz Cloud stocke et indexe ces embeddings. Lorsqu’une requête arrive, PyTorch intègre les données de la requête, et Zilliz Cloud effectue une recherche de similarité pour trouver les éléments correspondants les plus proches, permettant des applications telles que la recherche visuelle, les recommandations et la détection d’anomalies.

1. Installer les packages requis

pip install pymilvus torch gdown torchvision tqdm

2. Télécharger le jeu de données

Utilisez gdown pour télécharger le jeu de données Impressionist-Classifier et l’extraire :

import gdown
import zipfile

url = 'https://drive.google.com/uc?id=1OYDHLEy992qu5C4C8HV5uDIkOWRTAR1_'
output = './paintings.zip'
gdown.download(url, output)

with zipfile.ZipFile("./paintings.zip","r") as zip_ref:
    zip_ref.extractall("./paintings")

3. Définir les paramètres globaux

COLLECTION_NAME = 'image_search'
DIMENSION = 2048
MILVUS_HOST = "localhost"
MILVUS_PORT = "19530"
BATCH_SIZE = 128
TOP_K = 3

4. Configurer la collection Milvus

Connectez-vous à Milvus, créez une collection avec le schéma approprié et construisez un index :

from pymilvus import connections, utility, FieldSchema, CollectionSchema, DataType, Collection

connections.connect(host=MILVUS_HOST, port=MILVUS_PORT)

if utility.has_collection(COLLECTION_NAME):
    utility.drop_collection(COLLECTION_NAME)

fields = [
    FieldSchema(name='id', dtype=DataType.INT64, is_primary=True, auto_id=True),
    FieldSchema(name='filepath', dtype=DataType.VARCHAR, max_length=200),
    FieldSchema(name='image_embedding', dtype=DataType.FLOAT_VECTOR, dim=DIMENSION)
]
schema = CollectionSchema(fields=fields)
collection = Collection(name=COLLECTION_NAME, schema=schema)

index_params = {
    'metric_type':'L2',
    'index_type':"IVF_FLAT",
    'params':{'nlist': 16384}
}
collection.create_index(field_name="image_embedding", index_params=index_params)
collection.load()

5. Générer les embeddings et insérer les données

Chargez les images, générez les embeddings à l’aide du modèle ResNet50 pré-entraîné et insérez-les dans Milvus :

import glob
import torch
from torchvision import transforms
from PIL import Image
from tqdm import tqdm

paths = glob.glob('./paintings/paintings/**/*.jpg', recursive=True)

model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet50', pretrained=True)
model = torch.nn.Sequential(*(list(model.children())[:-1]))
model.eval()

preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

def embed(data):
    with torch.no_grad():
        output = model(torch.stack(data[0])).squeeze()
        collection.insert([data[1], output.tolist()])

data_batch = [[],[]]

for path in tqdm(paths):
    im = Image.open(path).convert('RGB')
    data_batch[0].append(preprocess(im))
    data_batch[1].append(path)
    if len(data_batch[0]) % BATCH_SIZE == 0:
        embed(data_batch)
        data_batch = [[],[]]

if len(data_batch[0]) != 0:
    embed(data_batch)

collection.flush()

6. Effectuer une recherche de similarité d’images

Intégrez les images de test et recherchez les peintures les plus similaires dans la collection :

import time
from matplotlib import pyplot as plt

search_paths = glob.glob('./paintings/test_paintings/**/*.jpg', recursive=True)

def embed(data):
    with torch.no_grad():
        ret = model(torch.stack(data))
        if len(ret) > 1:
            return ret.squeeze().tolist()
        else:
            return torch.flatten(ret, start_dim=1).tolist()

data_batch = [[],[]]

for path in search_paths:
    im = Image.open(path).convert('RGB')
    data_batch[0].append(preprocess(im))
    data_batch[1].append(path)

embeds = embed(data_batch[0])
start = time.time()
res = collection.search(embeds, anns_field='image_embedding', param={'nprobe': 128}, limit=TOP_K, output_fields=['filepath'])
finish = time.time()

f, axarr = plt.subplots(len(data_batch[1]), TOP_K + 1, figsize=(20, 10), squeeze=False)

for hits_i, hits in enumerate(res):
    axarr[hits_i][0].imshow(Image.open(data_batch[1][hits_i]))
    axarr[hits_i][0].set_axis_off()
    axarr[hits_i][0].set_title('Search Time: ' + str(finish - start))
    for hit_i, hit in enumerate(hits):
        axarr[hits_i][hit_i + 1].imshow(Image.open(hit.entity.get('filepath')))
        axarr[hits_i][hit_i + 1].set_axis_off()
        axarr[hits_i][hit_i + 1].set_title('Distance: ' + str(hit.distance))

plt.savefig('search_result.png')

• Recherche d’images avec PyTorch et Milvus — Tutoriel officiel de Milvus pour créer une recherche d’images avec PyTorch
• Recherche d’images avec Zilliz Cloud et PyTorch — Documentation de Zilliz Cloud pour l’intégration de PyTorch
• Partenariat de Zilliz avec PyTorch — Page de partenariat de Zilliz et PyTorch avec des tutoriels
• Améliorer l’expérience utilisateur avec des recommandations de mode basées sur l’image — Blog de Zilliz sur les recommandations basées sur l’image
• Bien démarrer avec la recherche hybride avec Milvus — Blog de Zilliz sur les techniques de recherche hybride