Оптимизация потребления памяти в ML-библиотеке LANCETNIC

Немного о библиотеке LANCETNIC

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

Вот GitHub репозиторий для тех кто хочет подробнее ознакомиться:

👉 https://github.com/Lancet52/lancetnic

Но вернемся к тому о чем хочу написать.

Недавно столкнулся с проблемой: моя библиотека lancetnic при обучении на больших текстовых датасетах просто перегружала оперативную память. На ноуте в 16 ГБ RAM модель не могла обучиться даже на 25 тыс. строк. Разбирался. Я начал разбираться и нашёл пару причин критического перерасхода памяти.

Причина № 1. Плотные матрицы вместо разряженных

TfidfVectorizer (который используется для векторизации текстовых данных) из sklearn по умолчанию возвращает разреженную матрицу (sparse). Он хранит только ненулевые значения.

TfidfVectorizer обрабатывает весь датасет и собирает все уникальные слова которые встречаются хотя бы в одном сообщении. Формируется словарь.

Для наглядности возьмём датасет спам сообщений:

Векторайзер собирает уникальные слова из всего датасета: мастер, маникюра, обучим, от, 7000, в, день, куплю, iphone, 15, недорого, срочно, бесплатный, кредит, без, справок, за, 1 час, пиши, личные, сообщения и т.д. Из всего этого формируется словарь в виде таблицы, где каждое сообщение превращается в строку из большого количества чисел (по одному числу из словаря):

Как видно, если слова нет в строке то столбец заполняется нулём. Чем больше датасет, тем больше слов и тем больше таблица (словарь).

ТехническиTfidfVectorizer создаёт матрицу которая хранит только ненулевые значения (это называется sparse или разреженная матрица).

В моем коде (в версии lancetnic 4.0.0) при векторизации я создавал полноценную матрицу с нулями и выглядело это выглядело так:

# Векторизация тектовых данныхdef vectorize_text(text_column, df_train, max_features):    if isinstance(text_column, str):        text_column = [text_column]    text_encoder_list = []    vectorizers = []            for text_col in text_column:        vectorizer_text = TfidfVectorizer(max_features=max_features)        # Вот тут        text_encoder = vectorizer_text.fit_transform(df_train[text_col].fillna('')).toarray()        text_encoder_list.append(text_encoder)        vectorizers.append(vectorizer_text)    # И вот тут    combined_text = sp.hstack(text_encoder_list).tocsr()    return combined_text, vectorizers

И если на датасете в 5000 строк это не выглядело критично для ноутбука, то на 25000 строк это съедало всю оперативную память. Поэтому решение этой проблемы было простым: я убрал .toarray() и оставил данные в sparse-формате:

# Векторизация тектовых данныхdef vectorize_text(text_column, df_train, max_features):    if isinstance(text_column, str):        text_column = [text_column]    text_encoder_list = []    vectorizers = []    for text_col in text_column:        vectorizer_text = TfidfVectorizer(max_features=max_features)        # Теперь так        text_encoder = vectorizer_text.fit_transform(df_train[text_col].fillna(''))         text_encoder_list.append(text_encoder)        vectorizers.append(vectorizer_text)    # И вот так    combined_text = sp.hstack(text_encoder_list).tocsr()    return combined_text, vectorizers

Ну и далее по коду все поправил.

Причина № 2. Двойное хранение данных в PyTorch Dataset

Класс Dataset при создании сразу конвертировал весь массив X в PyTorch-тензор через torch.tensor(X).
На примере кода для классификации:

class ClassifierDataset(Dataset):    def __init__(self, X, y):        # Создание копии в памяти        self.X = torch.tensor(X, dtype=torch.float32)        self.y = torch.tensor(y, dtype=torch.long)    def __len__(self):        return len(self.X)    def __getitem__(self, idx):        return self.X[idx], self.y[idx]

В конструктор класса Dataset передавалась матрица после TF-IDF, да к тому же еще и в полном виде (см. Причину № 1). Строкой self.X = torch.tensor(X, dtype=torch.float32) все это превращалось в Pytorch тензор. Это создавало вторую копию данных в памяти. Получалось, что данные хранились дважды.
Новый код:

class ClassifierDataset(Dataset):    def __init__(self, X, y):        self.is_sparse = sparse.issparse(X)        self.X = X        if not self.is_sparse:            self.X = torch.tensor(X, dtype=torch.float32)        self.y = torch.tensor(y, dtype=torch.long)      def __len__(self):        return self.X.shape[0]     def __getitem__(self, idx):        if self.is_sparse:            x = torch.tensor(self.X[idx].toarray(), dtype=torch.float32).squeeze(0)        else:            x = self.X[idx]                    return x, self.y[idx]

В конструкторе класса теперь остается только ссылка на массив из TF-IDF. А в getitem я беру только одну строку из матрицы и превращаю ее в тензор. Теперь в памяти в каждый момент времени находится только один батч, а не весь датасет.

Резюмируя

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

Буду рад всевозможным отзывам и обратной связи. Заинтересован чтобы LANCETNIC прошёл обкатку на реальных кейсах. Парочка таких уже есть. И один из них это антиспамбот (TAB) для телеги. О нём я писал уже в своей статье на Хабре: ссылка на статью

Официальный сайт LANCETNIC: https://lancetnic.ru/
GitHub: https://github.com/Lancet52/lancetnic
Блог разработки в ТГ: https://t.me/markovstate
Дашборд антиспам бота TAB: https://tab.lancetnic.ru/