Странная ошибка “assert d == self.d” в FAISS с ResNet50

Вопрос или проблема

import faiss
import os
import numpy as np
import torch
from torchvision.models import resnet50
from torchvision.models import ResNet50_Weights
from torchvision import transforms
from PIL import Image

accepted_formats = [".jpg", ".png", ".jpeg"]

image_dir = "images"

preprocess = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor()
])

weights = ResNet50_Weights.DEFAULT

def load_images_into_index(index):
    images = os.listdir(image_dir)
    model = resnet50(weights=weights)
    model.eval()
    for image in images:
        if image.lower().endswith(tuple(accepted_formats)):
            image_tensor = preprocess(Image.open(image_dir + "/" + image).convert("RGB"))
            with torch.no_grad():
                features = model(image_tensor.unsqueeze(0))
            features = features.detach().cpu().numpy()
            faiss.normalize_L2(features)
            index.add(features)

index = faiss.IndexFlatL2()
load_images_into_index(index)

Ответ или решение

Ошибка assert d == self.d в FAISS возникает, когда размерность векторов, добавляемых в индекс, не совпадает с размерностью, которая уже была установлена для этого индекса. В вашем случае это может происходить из-за неправильной обработки выходных данных модели ResNet50, что приводит к добавлению векторов в индекс с неожиданной размерностью.

Для решения этой проблемы, пожалуйста, проверьте следующие шаги:

1. Изменение размерности выходных данных: Результаты, которые вы получаете из модели ResNet50, имеют размерность, соответствующую количеству классов (обычно 1000 для предобученной модели на ImageNet). Чтобы получить вектор признаков фиксированной размерности (например, 2048 для ResNet50), вам нужно использовать выход из предпоследнего слоя (это можно сделать с помощью отзыва слоя).

   Обновите код следующим образом:

   def load_images_into_index(index):
      images = os.listdir(image_dir)
      model = resnet50(weights=weights)
      model.eval()
   
      # Извлеките нужный слой для получения признаков
      model = torch.nn.Sequential(*(list(model.children())[:-1]))  # Убираем последний слой
   
      for image in images:
          if image.lower().endswith(tuple(accepted_formats)):
              image_tensor = preprocess(Image.open(image_dir + "/" + image).convert("RGB"))
              with torch.no_grad():
                  features = model(image_tensor.unsqueeze(0))
              features = features.view(features.size(0), -1)  # Преобразуем вектор в 1D
              features = features.detach().cpu().numpy()
              faiss.normalize_L2(features)
              index.add(features)
   
   index = faiss.IndexFlatL2(2048)  # Укажите размерность индекса
   load_images_into_index(index)

2. Проверка размерности после извлечения признаков: После изменения кода обязательно проверьте форму извлеченных признаков перед добавлением их в индекс:

   print(features.shape)  # Это должно выводить (1, 2048)

   Убедитесь, что размерность признаков соответствует d=2048, так как именно такая размерность будет требоваться для использования IndexFlatL2.

3. Установка правильной размерности для индекса: Убедитесь, что вы создаете индекс FAISS с правильной размерностью вектора. Вам нужно передать размерность 2048 во время инициализации индекса:

   index = faiss.IndexFlatL2(2048)

После выполнения всех этих шагов проблема с ошибкой assert d == self.d должна быть решена. Если ошибки продолжат появляться, внимательно просмотрите вывод, чтобы проверить, соответствуют ли размеры векторов в процессе извлечения признаков и добавления их в индекс.