joblib и mmap: загрузить массивы модели без полной копии

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

Отображение работает с массивами, а не со всем объектом

joblib.load с mmap_mode="r" восстанавливает подходящие массивы как NumPy memmap. Их данные читаются из отображённого файла, и операционная система может совместно использовать страницы между процессами. Python-объекты вокруг массивов и промежуточные результаты predict при этом не становятся общими автоматически.

Файл должен быть сохранён без сжатия. Сжатый поток нельзя непосредственно отобразить как массив с нужными смещениями. Поэтому выбор между меньшим размером на диске и возможностью mmap делают на этапе dump. Режим r запрещает изменять отображённые данные через такой массив.

Проверить конкретный обученный массив

В примере LinearRegression создаёт coef_. После сохранения и загрузки проверяем его реальный тип, признак writeable и прогноз. Небольшой массив нужен, чтобы показать механизм; численного выигрыша памяти такой размер не обещает. Для вашей модели сначала проверьте, какие её части действительно стали memmap.

Создаётся только доверенный локальный артефакт. joblib использует pickle-механизм, поэтому выбирать mmap вместо обычной загрузки не значит менять требования к происхождению файла. Загружайте модель из своей контролируемой сборки в совместимом окружении.

import gc
from pathlib import Path
from tempfile import TemporaryDirectory
import joblib
import numpy as np
from sklearn.linear_model import LinearRegression

model = LinearRegression().fit([[0], [1], [2], [3]], [1, 3, 5, 7])
with TemporaryDirectory() as directory:
    path = Path(directory) / "model.joblib"
    joblib.dump(model, path, compress=0)
    loaded = joblib.load(path, mmap_mode="r")
    assert isinstance(loaded.coef_, np.memmap)
    assert not loaded.coef_.flags.writeable
    predicted = loaded.predict([[4]])[0]
    assert abs(predicted - 9) < 1e-9
    print(type(loaded.coef_).__name__, float(predicted))
    del loaded
    gc.collect()

Сохранить файл на всё время работы

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

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

Закрепить базовую операцию

В открытом уроке сохранения через joblib вы пройдёте dump, load и прогноз восстановленной модели. Затем измените способ хранения и проверьте, какие массивы вашей модели поддерживают отображение. Откройте урок «Сохранение обученной модели через joblib». Затем вернитесь к сценарию статьи и проверьте его на своём небольшом наборе. Посмотрите программу курса первой ML-модели, чтобы пройти путь от обучения к использованию результата.