В этом посте мы экспортируем и квантуем визуальный кодировщик предварительно обученной модели CLIP (ViT-B/32), используя ONNX Runtime Web. Это квантование снизит точность весов модели, уменьшив размер сохраненного файла до 25% от его исходного размера и обеспечив значительное сокращение использования памяти. Благодаря уменьшенной модели и преимуществам, предоставляемым ONNX, мы надеемся увидеть возможность запуска таких больших и сложных моделей глубокого обучения, как CLIP, в JavaScript на стороне клиента.

Мы сравниваем производительность исходной и квантованной моделей при создании вложений признаков для модели логистической регрессии с использованием набора данных CIFAR100. Мы обнаружили, что квантованная модель не страдает в отношении точности проверки, но фактически немного превосходит исходную модель.

Хотя это не означает, что квантованная модель так же эффективна, как и исходная, в других приложениях, она обеспечивает доказательство концепции того, что квантованная модель имеет некоторую полезность. Мы рассмотрим концепции и код для воспроизведения этого эксперимента и регистрации его результатов в Comet.

Следуйте подробному коду через эту записную книжку Colab.

Введение

В Comet мы помогаем командам, занимающимся машинным обучением и наукой о данных, быстрее создавать более качественные модели, расширяя границы наблюдаемости машинного обучения.

С Comet Panels экспериментальные и производственные показатели у вас под рукой в ​​браузере.

В настоящее время мы изучаем новые способы улучшения наблюдаемости путем запуска моделей машинного обучения в браузере. Таким образом, мы были в восторге от ONNX Runtime, кроссплатформенной среды глубокого обучения, которая может значительно повысить скорость логического вывода на процессорах.

Мы изучаем возможности встраивания обученных моделей в панели Comet с использованием платформы ONNX, которая может включать в себя комбинацию готовых панелей, которые используют глубокое обучение, чтобы помочь вам ориентироваться в ваших проектах, или шаблоны для тестирования вашего собственного ONNX. модели в панелях.

Однако машинное обучение в браузере сопряжено со сложностями. Во-первых, это означает развертывание модели с гораздо меньшим контролем над оборудованием, на котором она будет использоваться. Часто это означает использование логического вывода ЦП, который все еще может быть чрезмерно медленным по сравнению с развертываниями на основе графического процессора — даже с увеличением производительности ONNX.

Более того, многие модели глубокого обучения имеют миллионы или миллиарды параметров, и эти параметры необходимо сохранять в относительно больших файлах. Нередко можно увидеть, как модель занимает сотни мегабайт памяти, если не больше.

С этой целью нас очень интересуют методы сжатия моделей ML для уменьшения объема памяти, файлов меньшего размера и более быстрого вывода. К ним относятся дистилляция знаний, обрезка и квантование.

В этом посте мы сосредоточимся на квантовании.

Что такое квантование после обучения?

Вкратце, квантование после обучения — это простая (но не всегда легкая) процедура сжатия обученной модели. Мы могли бы думать о квантовании как об округлении чисел, но правильнее было бы сказать, что мы «уменьшаем точность» чисел.

Если бы наша модель изначально хранила параметры как 64-битные значения с плавающей запятой, мы могли бы квантовать их в 32-битные значения с плавающей запятой или, возможно, 8-битные целые числа. Мы не переучиваем существующую модель и не обучаем новую, урезанную модель. Мы просто конвертируем значения в менее точный тип данных.

Это уменьшает размер модели на диске и в памяти и ускоряет вывод, но обычно это требует затрат. Точно так же, как при сжатии изображения, мы отбрасываем некоторую (или большую часть) нетривиальную информацию. Часто это приводит к снижению производительности модели.

Существуют более мощные методы сжатия моделей ML, и мы также изучаем их для CLIP. Например, есть очень интересная работа, которую делает Neural Magic. Но квантование после обучения послужит для нас недорогой базой.

Что такое КЛИП?

В январе 2020 года независимая исследовательская группа OpenAI опубликовала свои первые результаты по Contrastive Language-Image Pre-training (CLIP), методу, который одновременно изучает сопоставимые кодировки текстов и изображений путем сравнения изображений с их подписями. Таким образом, когда мы говорим о модели CLIP, мы говорим о двух обученных моделях: визуальном кодировщике и текстовом кодировщике. Как может выглядеть это кодирование и сравнение?

Допустим, у нас есть две картинки. На одной изображена собака, а на другой кошка. Мы запускаем каждый из них через визуальный кодировщик, чтобы получить встраивание. В качестве аргумента предположим, что это вложение представляет собой вектор формы (1, 512) (типично для предварительно обученных моделей CLIP). У нас также есть несколько подписей — скажем, «изображение собаки» и «фото кошки». После токенизации текста мы можем пропустить эти подписи через кодировщик текста и получить еще два встраивания размера (1, 512).

Теперь мы можем сравнивать текстовую и визуальную кодировки на одном уровне! Мы могли бы использовать что-то вроде косинусного сходства, чтобы вычислить сходство между изображением собаки и каждой подписью. Если в дрессировке все прошло по плану, следует ожидать, что фото собаки будет лучше сочетаться с подписью «фото собаки», чем с «фото кота». А поскольку предварительно обученные модели OpenAI обучились на 40 миллиардах пар изображений и подписей, мы можем с полным основанием ожидать, что они видели некоторые фотографии животных.

Экспертные модели, такие как CLIP, очень дороги и трудоемки в обучении. Мы получаем от них ценность, расширяя их для применения к широкому кругу задач. На самом деле, адаптация экспертных моделей к новым задачам или данным является активной областью исследований в области машинного обучения.

Мы могли бы использовать CLIP для нулевой классификации, используя оценки сходства между кодировками изображения и текста. В качестве альтернативы мы могли бы использовать CLIP для семантического поиска изображений; мы вводим описание изображения, пропускаем его через текстовый кодировщик CLIP, а затем сравниваем полученное вложение с набором изображений.

В качестве альтернативы, тем не менее, мы могли бы использовать CLIP для генеративных задач, таких как объединение его с GAN или другим механизмом генерации изображений и оптимизация изображения для соответствия конкретной текстовой подсказке.

Сравнительный анализ производительности CLIP

Как узнать, хорошо ли работает модель CLIP? Один из вариантов — использовать его для простой задачи классификации с широко изученным набором данных вне выборки. Это помогает определить надежность модели, насколько хорошо она обобщает новые, но похожие данные.

У нас есть возможность выполнить нулевую классификацию, в которой мы будем кодировать исходные изображения, а также метки (т. е. подписи). Или мы могли бы выделить визуальный кодировщик и использовать все наши сгенерированные вложения в качестве новых функций для простой модели классификации, такой как логистическая регрессия.

OpenAI включает пример этого в репозиторий и документ для CLIP, в котором визуальный кодировщик используется для предварительной обработки вложений в наборе данных CIFAR-100. CIFAR-100 назван в честь включения 100 различных классов изображений; есть еще попроще СИФАР-10.

В этом эксперименте мы собираемся сравнить исходные предварительно обученные модели CLIP с квантованными с использованием теста «линейный зонд» CIFAR-100. Мы надеемся увидеть аналогичную прогностическую эффективность в квантованной модели, несмотря на ее меньший размер и пониженную точность ее параметров.

Теперь, если квантованная модель работает так же хорошо, как оригинал в этом тесте, это не означает, что квантованная модель «так же хороша», как оригинал в целостном смысле. Но это означало бы, что квантованная модель не дает сбоев в процессе квантования. Давайте углубимся в код.

Пытаетесь отследить и воспроизвести сложные параметры эксперимента? Артефакты — это лишь один из многих инструментов в наборе инструментов Comet, упрощающих управление моделями. Прочитайте наш пример использования PetCam, чтобы узнать больше.

Загрузить модель CLIP

Доступ к архитектуре модели CLIP и предварительно обученным весам можно получить через библиотеку clip, которую можно установить из этого репозитория. Подробнее о работе с моделями CLIP читайте в моем предыдущем посте:



Но есть несколько вещей, которые я хочу упомянуть о следующем коде:

import clip
device = 'cpu'
model_name = 'ViT-B/32'
model, preprocess = clip.load(model_name, device)

Во-первых, обратите внимание, что мы установили устройство PyTorch на 'cpu'. Поскольку наша модель ONNX будет работать на ЦП, мы часто хотим установить модель PyTorch на ЦП. Если мы попытаемся отследить вычислительный граф на графическом процессоре, мы можем столкнуться с некоторыми ошибками или неожиданными результатами.

Нам нужно было указать конкретную конфигурацию для CLIP, модель 'ViT-B/32'. Это читается как Преобразователь зрения (базовый патч 32). Итак, это модель Transformer (чтобы понять, что это значит, прочтите эту запись в блоге), которая разбивает изображение на патчи размером 32x32. У CLIP также есть другая модель ViT, ViT-B/16, которая имеет базовые патчи размером 16x16. Эта модель, конечно, сложнее, потому что маленькие патчи = больше патчей.

В библиотеке clip также есть предварительно обученные модели с архитектурами на основе ResNet. Я говорю об этом, потому что код квантования в этом посте работает с "ViT-B/32", но не с другими моделями в настоящее время. В частности, мы все еще работаем над реализацией "ViT-B/16".

Наконец, обратите внимание, что clip.load дает нам модель, а также функцию для предварительной обработки изображений. Помимо прочего, это преобразует изображение из PIL в тензор PyTorch, соответствующим образом изменяет его размер (эти модели работают с изображениями RGB 224x224) и нормализует цветовые каналы. Этот этап предварительной обработки отделен от самой модели и не будет захвачен при экспорте ONNX.

Поэтому в будущем, когда мы захотим запустить эту модель в веб-браузере с клиентским JavaScript, нам нужно будет воспроизвести этот этап предварительной обработки отдельно с библиотеками JavaScript. Сегодня мы просто хотим посмотреть, жизнеспособна ли квантовая модель, так что мы не будем об этом беспокоиться.

Экспорт ONNX

PyTorch имеет встроенную поддержку для работы с ONNX, а модели можно экспортировать с помощью torch.onnx.export, как показано ниже. Теперь «моделью» в этом случае обычно может быть модуль PyTorch, но не обязательно. Мы также можем экспортировать некоторые функции и скрипты в ONNX, но сейчас самое время упомянуть, что путь от PyTorch к ONNX не всегда гладок. Существуют ограничения на то, что можно экспортировать и как. Во-первых, существует проблема отслеживания vs. сценарий. У нас нет места для слишком подробного описания этого, поэтому я буду краток и постараюсь не делать здесь грубых ошибок.

Python — это интерпретируемый язык. Это означает, что сценарии, которые мы пишем, считываются при каждом запуске. Во многих случаях это неэффективно по сравнению с компилируемыми языками, такими как C или Java. Многие высокопроизводительные библиотеки Python (ярким примером является NumPy) используют функции Python для вызова более оптимизированных подпрограмм на скомпилированном языке, что обеспечивает значительное ускорение.

Такие библиотеки, как PyTorch и TensorFlow, делают нечто подобное. Они компилируют скрипт в график вычислений, что-то вроде функциональной блок-схемы операций.

Однако PyTorch не может скомпилировать любой произвольный код в граф. Некоторые вещи, например операторы if-else в чистом Python, не отображаются на графике. Одним из вариантов является написание кода с использованием TorchScript, который является подмножеством Python, который PyTorch может легко скомпилировать. Другой вариант — позволить PyTorch отследить график. По сути, это записывает все, что происходит при вызове функции/модели PyTorch, и компилирует график на основе этого.

Однако трассировка не сможет зафиксировать такие вещи, как чистый поток управления Python (операторы if-else), и вместо этого всегда будет работать так же, как при первоначальной трассировке графа. Более подробное объяснение смотрите в документации. Суть в том, что мы теряем несколько вещей при наивном подходе к отслеживанию.

Все это говорит о том, что метод, который мы показываем для квантования CLIP, не является лучшим способом сделать это. Во-первых, у нас нет переменных размеров партий. Часть нашей дальнейшей работы над этим проектом будет включать пересмотр этого процесса. Но пока это работает достаточно хорошо.

import torch
# Save model to onnx
  
torch.onnx.export(model.visual,                        
                  input_tensor,       
                  filepath,                     
                  export_params=True,           
                  opset_version=11,            
                  do_constant_folding=True,    
                  input_names = ['input'],   
                  output_names = ['output'])

Давайте немного посмотрим на этот код:

  • Первый аргумент — визуальный кодировщик из предварительно обученной модели CLIP. Это объект torch.nn.Module.
  • Далее мы передаем входной тензор. В данном случае это будет использоваться для выполнения трассировки графа вычислений модели. Таким образом, этот входной тензор должен иметь подходящую форму и тип для использования в модели. Подробнее об этом в блокноте Colab или в моем вступительном посте о работе с CLIP.
  • Далее мы передаем путь к файлу, в который мы можем сохранить модель ONNX.
  • Далее мы указываем, что хотим экспортировать подобранные параметры модели вместе с графиком. В противном случае мы не можем сделать вывод!
  • opset_version: здесь описывается, какие операторы ONNX должны соответствовать операциям в torch. Здесь я снова укажу на документацию PyTorch для получения дополнительной информации.
  • do_constant_folding: Где это возможно, это объединяет (или «сворачивает») последовательности операций в одну операцию, что может значительно ускорить время вывода.
  • input_names: список имен, которые мы выбираем для входных данных графика. В этом случае наша модель имеет один вход, и мы можем назвать ее как угодно. Это имя понадобится нам позже, когда мы попытаемся сделать выводы в ONNX Runtime.
  • output_names: Аналогично, имя или имена для выходных тензоров.

Квантование

После некоторых проб и ошибок я остановился на подходе, заключающемся в экспорте неквантованной модели в ONNX, а затем в использовании ONNX Runtime для квантования. PyTorch также имеет возможности квантования, и теоретически они должны быть в значительной степени совместимы с ONNX, но я обнаружил, что в этом случае это работает лучше.

Другая причина, по которой я позволил ONNX выполнять квантование, заключается в том, что я также хотел проверить производительность не квантованной модели ONNX CLIP. Он должен быть почти идентичен исходной модели факела с точки зрения прогнозов, поэтому мы будем следить за этим, когда приступим к эксперименту. Если мы не наблюдаем идентичных или почти идентичных кодировок между PyTorch и неквантованными моделями ONNX, значит, что-то пошло не так.

from onnxruntime.quantization import quantize_dynamic, QuantType
quantized_model = quantize_dynamic(input_filepath,
                                   output_filepath,
                                   weight_type=QuantType.QUInt8)

Важно отметить здесь аргумент weight_type=QuantType.QUInt8. Мы решили квантовать эту модель до 8-битных целых чисел. Первоначально мы имели дело с 32-битными значениями с плавающей запятой для параметров модели. Это объясняет, почему у квантованной модели размер файла примерно на 25 % больше, чем у оригинала — для хранения каждого параметра требуется на 25 % больше битов!

Создание вложений

Как я упоминал ранее, здесь мы используем классический набор данных CIFAR100, который, будучи классическим, встроен в torchvision.datasets. Мы можем настроить наши объекты Dataset следующим образом:

from torch.utils.data import DataLoader
from torchvision.datasets import CIFAR100
from tqdm import tqdm
# Load the dataset
root = os.path.expanduser("~/.cache")
train = CIFAR100(root, download=True, train=True, transform=preprocess)
test = CIFAR100(root, download=True, train=False, transform=preprocess)

Теперь к train и test можно получить доступ несколькими способами. Типичный идиоматический способ сделать это в PyTorch — использовать DataLoader, который устанавливает итератор, который мы можем использовать для получения точек данных по мере необходимости. Это может выглядеть примерно так:

for image, label in tqdm(DataLoader(train, batch_size=10)):
	# Do stuff here
	...

Обратите внимание, что мы используем популярную библиотеку tqdm, чтобы получить хороший индикатор выполнения. Вам может быть интересно, для какой цели служит DataLoader. Среди прочего, это дает нам контроль над процессом предварительной обработки... процесса.

Например, мы можем указать количество работ, которые необходимо использовать для предварительной обработки необработанных объектов (сама функция предварительной обработки была передана в аргумент transform, когда мы инициализировали наборы данных CIFAR100 выше). Вообще говоря, эффективное использование DataLoaders может помочь устранить узкие места ЦП в глубоком обучении.

Чтобы получить дополнительную информацию о таких узких местах, ознакомьтесь с другим моим сообщением о Heartbeat.



Теперь давайте посмотрим, как мы генерируем наши вложения в PyTorch по сравнению с ONNX, начиная с PyTorch.

def get_features_torch(dataset, batch_size=16):
    all_features = []
    all_labels = []
    with torch.no_grad():
        for images, labels in tqdm(DataLoader(dataset, batch_size=16)):
            features = model.encode_image(images.to(device))
            all_features.append(features)
            all_labels.append(labels)
    return torch.cat(all_features).cpu().numpy(), torch.cat(all_labels).cpu().numpy()

Несколько вещей, на которые следует обратить внимание: во-первых, я установил размер пакета равным 16. В блокноте Colab мы используем вывод графического процессора для исходной модели факела, поэтому мы выигрываем, делая пакетные прогнозы параллельно.

Я должен подчеркнуть, что даже с ускорением, обеспечиваемым ONNX Runtime и квантованием, исходная модель на хорошем графическом процессоре намного быстрее, чем наши модели ONNX, обрабатывая все 60 000 изображений за минуты, в то время как модели ONNX занимают часы. Но давайте помнить, что процессоры, как правило, дешевы и широко доступны по сравнению с современными графическими процессорами, которые мы получаем с Colab.

В конечном счете, почти все, что нам нужно сделать с этой функцией, — это перебрать набор данных и использовать model.encode_image для создания наших вложений, а затем сложить все это вместе в конце. Теперь давайте посмотрим на код ONNX:

def get_features_onnx(dataset, onnx_model_path):
    all_features = []
    all_labels = []
    
    ort_session = onnxruntime.InferenceSession(onnx_model_path)
    for image, label in tqdm(DataLoader(dataset, batch_size=1)):
        ort_inputs = {'input': to_numpy(image)} 
        features = ort_session.run(None, ort_inputs)[0]
        all_features.append(features)
        all_labels.append(label)
    return np.vstack(all_features), np.array(all_labels)

Что бросается в глаза по сравнению с кодом факела? Ну, во-первых, мы должны создать этот объект onnxruntime.InferenceSession, который загружает сохраненный график и готовит его для вывода. Мы также должны настроить словарь входных данных. Наконец, мы можем вызвать ort_session.run, чтобы превратить наши входы в выходы. Напомним, что мы создали наш график с одним выходом (называемым «выход»). Функция ort.session_run вернет список выходных данных, но, поскольку мы знаем, что у нас есть только один выходной сигнал, я индексирую этот список и просто беру в нем первое (наш единственный выходной сигнал).

Запись вещей на потом

Модели ONNX медленнее по сравнению с запуском модели PyTorch на высокопроизводительном графическом процессоре, поэтому мы не хотим снова и снова вычислять вложения. Мы могли бы захотеть использовать их позже для других задач, таких как другие дальнейшие анализы. Артефакты комет позволяют нам регистрировать, отслеживать, версии и повторно использовать большие коллекции активов, таких как наборы данных.



Мы также можем регистрировать наши модели ONNX с помощью Comet, что избавляет нас от необходимости экспортировать их позже (и при этом улучшает воспроизводимость!). Кроме того, мы могли бы также рассмотреть возможность регистрации этих зарегистрированных моделей в Реестре моделей Comet, который предназначен для интеграции с производственными развертываниями. Мы пока не будем этого делать, а логируем модели.

См. упрощенную версию кода блокнота ниже:

# Create an experiment to log results to
e = comet_ml.Experiment()
# Log the model
e.log_model(f'clip-visual-quant', path_to_saved_model)
# Create a new artifact
artifact = comet_ml.Artifact(f"clip-cifar100-quant", "dataset")
# Generate features
X_train, y_train = get_features_onnx(train, fullmodel_paths[model_dir])
X_test, y_test = get_features_onnx(test, fullmodel_paths[model_dir])
# Serialize the embeddings to files
np.save(f'{feature_dir}/{model_dir}/X_train.npy', X_train)
np.save(f'{feature_dir}/{model_dir}/X_test.npy', X_test)
np.save(f'{feature_dir}/{model_dir}/y_train.npy', y_train)
np.save(f'{feature_dir}/{model_dir}/y_test.npy', y_test)
  
# Add the embedding files to the artifact
artifact.add(f'{feature_dir}/{model_dir}/X_train.npy')
artifact.add(f'{feature_dir}/{model_dir}/X_test.npy')
artifact.add(f'{feature_dir}/{model_dir}/y_train.npy')
artifact.add(f'{feature_dir}/{model_dir}/y_test.npy')
e.log_artifact(artifact)
e.end()

Еще одним преимуществом использования Артефактов является то, что мы можем сделать их общедоступными! В блокноте Colab вы можете пропустить этот шаг и использовать Comet для загрузки встроенных функций, которые я уже создал. Для получения дополнительной информации об артефактах комет см. следующую документацию.

Теперь, когда у нас есть функции, мы готовы провести эксперимент по классификации.

Логистическая регрессия

Полный код этого эксперимента можно найти в блокноте Colab — в основном это базовое обучение scikit-learn, поэтому я не буду подробно его описывать, но по сути мы делаем что-то вроде этого:

from sklearn.linear_model import LogisticRegression
classifier = LogisticRegression()
classifier.fit(X_train, y_train)
predictions = classifier.predict(X_test)

Где X_train и X_test — это встраивание функций для конкретной модели, а y_train и y_test — правильные метки классов. Оттуда мы вычисляем некоторые метрики обучения/проверки и записываем все в Comet для сравнения.

Опять же, поскольку этот пост больше посвящен экспорту в ONNX и квантизации с помощью ONNX, я не буду подробно рассматривать этот код. Полный подробный пример использования Comet для классификации можно найти в прекрасном посте моего коллеги Harpreet Sahota здесь:



Полученные результаты

В конечном итоге для каждой модели мы находим следующее: модели PyTorch и ONNX имеют почти одинаковую точность проверочного набора 0,7964 и 0,7967 соответственно, в то время как квантованная модель достигает точности 0,805.

Это простая метрика для оценки простого эксперимента по сравнению очень и очень сложных моделей. Но в узких рамках этого эксперимента мы обнаружили, что простое квантование визуального кодировщика CLIP не полностью ломает ситуацию. Это открывает несколько следующих шагов:

  • Квантование текстового кодировщика: модель CLIP содержит визуальный кодировщик и текстовый кодировщик, и до сих пор мы рассматривали только визуальный кодировщик. Чтобы вычислить скрытые представления ввода текста, нам нужно обработать экспорт этого текстового кодировщика и некоторых дополнительных функций обработки текста в ONNX.
  • Классификация нулевого выстрела. В этом методе используется кодировщик текста. Сначала мы создадим вложение, соответствующее каждой метке класса, например. f"A picture of {class_name}" для каждого класса в наших данных. Затем, чтобы классифицировать изображение, мы кодируем его и сравниваем встраивание с каждым отдельным вложением текста. Прогнозируемый класс — это тот, чье встраивание текста больше всего похоже на встраивание изображения.

В следующем посте этой серии мы рассмотрим квантизацию/экспорт текстового кодировщика и нулевой классификации, а также то, как мы можем начать регистрировать/визуализировать артефакты модели ONNX с помощью Comet.

Примечание редактора. Heartbeat — это интернет-издание и сообщество, созданное участниками и посвященное предоставлению лучших образовательных ресурсов для специалистов по науке о данных, машинному обучению и глубокому обучению. Мы стремимся поддерживать и вдохновлять разработчиков и инженеров из всех слоев общества.

Независимая от редакции, Heartbeat спонсируется и публикуется Comet, платформой MLOps, которая позволяет специалистам по данным и командам машинного обучения отслеживать, сравнивать, объяснять и оптимизировать свои эксперименты. Мы платим нашим авторам и не продаем рекламу.

Если вы хотите внести свой вклад, перейдите к нашему призыву к участию. Вы также можете подписаться на получение наших еженедельных информационных бюллетеней (Еженедельник глубокого обучения и Информационный бюллетень Comet), присоединиться к нам в Slack и следить за Comet в Twitter и LinkedIn, чтобы получать ресурсы, события и многое другое, что поможет вам быстрее создавать лучшие модели машинного обучения.