Неофициальная c++ LibTorch имплементация RuCLIP (SberAI)

Подробности: Категория: Проекты; Просмотров: 682

RuCLIP (Russian Contrastive Language–Image Pretraining) - это мультимодальная нейросеть (текст, изображение) обученная находить соответствия изображений их текстовым описаниям и наоборот. Оригинальный CLIP был разработан в OpenAI и поддерживает английский язык. https://arxiv.org/pdf/2103.00020.pdf. В компании SberAI обучили модель для русского языка и вложили веса и PyTorch код которые теперь можно найти тут: https://github.com/ai-forever/ru-clip. Ну а я, в свою очередь, перенес на c++ и LibTorch. RuCLIP в основе визуального и текстового кодировщиков имеет один и тот же трансформер поэтому модель получилась очень лаконичная и простая в использовании. Используется эффективный токенайзер от vk.com: YouTokenToMe tokenizer https://github.com/VKCOM/YouTokenToMe который написан на c++ и поэтому мне было легко его использовать.

Возьмем следующие изображения:

и следующие текстовые метки:
{"кот", "медведь", "лиса"}
и посчитаем матрицу вероятностей соответствия этих изображений этим меткам используя самую большую модель на 475млн параметров!

#include "TorchHeader.h"
 
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/imgproc/types_c.h>
#include <opencv2/highgui/highgui.hpp>
 
#include "RuCLIP.h"
#include "RuCLIPProcessor.h"
 
int main(int argc, const char* argv[])
{
 setlocale(LC_ALL, "");
 const int INPUT_IMG_SIZE = 336;
 torch::manual_seed(24);
 
 torch::Device device(torch::kCPU);
 if (torch::cuda::is_available())
 {
  std::cout << "CUDA is available! Running on GPU." << std::endl;
  device = torch::Device(torch::kCUDA);
 } else {
  std::cout << "CUDA is not available! Running on CPU." << std::endl;
 }
 
 CLIP clip = FromPretrained("..//data//ruclip-vit-large-patch14-336");
 clip->to(device);
 
 RuCLIPProcessor processor(
  "..//data//ruclip-vit-large-patch14-336//bpe.model",
  INPUT_IMG_SIZE,
  77,
  { 0.48145466, 0.4578275, 0.40821073 },
  { 0.26862954, 0.26130258, 0.27577711 }
 );
 
  ////Или можно без него сначала попробовать
  //RuCLIPPredictor(clip, processor, device, templates, 8);
 
 //Загрузить картинки
 std::vector <cv::Mat> images;
 images.push_back(cv::imread("..//data//test_images//1.png", cv::ImreadModes::IMREAD_COLOR));
 images.push_back(cv::imread("..//data//test_images//2.jpg", cv::ImreadModes::IMREAD_COLOR));
 images.push_back(cv::imread("..//data//test_images//3.jpg", cv::ImreadModes::IMREAD_COLOR));
 //resize->[336, 336]
 for (auto &it : images)
   cv::resize(it, it, cv::Size(INPUT_IMG_SIZE, INPUT_IMG_SIZE));
 
 //Завести метки
 std::vector<std::string> labels;
 labels = {"кот", "медведь", "лиса"};
 
 auto dummy_input = processor(labels, images);
 try {
  torch::Tensor logits_per_image = clip->forward(dummy_input.first.to(device), dummy_input.second.to(device));
  torch::Tensor logits_per_text = logits_per_image.t();
  auto probs = logits_per_image.softmax(/*dim = */-1).detach().cpu();
  std::cout << "probs per image: " << probs << std::endl;
 } catch (std::exception& e) {
  std::cout << e.what() << std::endl;
 }
}

Получим следующий вывод:
0.8879 0.0063 0.1058
0.0014 0.0026 0.9960
0.0002 0.9994 0.0003
Очевидно, шалость удалась!

Вот ссылка на репозиторий с проектом: https://github.com/DeliriumV01D/RuCLIP