Оптимизация потребления памяти в 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/