Пример использования автоблока

Вычислительные графы в bosk требуют, чтобы их узлы были экземплярами BaseBlock. Чтобы упростить включение пользовательских моделей, был разработан декоратор auto_block. Используя его, вы можете легко преобразовать свою пользовательскую модель в BaseBlock. Чтобы получить примеры преобразования моделей интерфейса scikit-learn, см. код подпакета bosk.block.zoo.models.classification.

Для выполнения этого примера у Вас должен быть установлен Pytorch. Если у вас не установлен torch пакет в Python, выполните следующую секцию:

!pip install torch

from bosk.block.auto import auto_block
from bosk.block.meta import BlockExecutionProperties
from bosk.executor.parallel.greedy import GreedyParallelExecutor
from bosk.painter.graphviz import GraphvizPainter
from bosk.pipeline.builder.functional import FunctionalPipelineBuilder
from bosk.stages import Stage
import torch
from torch.utils.data import TensorDataset, DataLoader
import numpy as np
from sklearn.datasets import make_moons
from sklearn.model_selection import train_test_split
from IPython.display import Image

Определим следующую модель:

@auto_block(
    execution_props=BlockExecutionProperties(cpu=False, gpu=True, plain=False, threadsafe=True),
    random_state_field=None,
)
class NeuralNetwork(torch.nn.Module):
    def __init__(self, dim: int) -> None:
        super().__init__()
        self.device = torch.device('cuda')
        self.nn = torch.nn.Sequential(
            torch.nn.Linear(dim, dim * 10),
            torch.nn.ReLU(),
            torch.nn.Linear(dim * 10, dim),
            torch.nn.Tanh(),
            torch.nn.Linear(dim, 1),
            torch.nn.Sigmoid()
        ).to(self.device)
        self.optimizer = torch.optim.Adam(self.nn.parameters(), 0.0005)
        self.batch_size = 64
        self.epochs_num = 500
        self.loss_fn = torch.nn.CrossEntropyLoss()

    def _get_tensor(self, x: np.ndarray) -> torch.Tensor:
        return torch.from_numpy(x.astype(np.float32)).to(self.device)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        proba = self.nn(x)
        return torch.concat((1 - proba, proba), dim=-1)

    def transform(self, x: np.ndarray) -> np.ndarray:
        with torch.no_grad():
            x = self._get_tensor(x)
            return self(x).cpu().numpy()

    def fit(self, x: np.ndarray, y: np.ndarray) -> 'NeuralNetwork':
        x_tens = self._get_tensor(x)
        y_tens = self._get_tensor(y).long()
        dataset = TensorDataset(x_tens, y_tens)
        data_loader = DataLoader(dataset, self.batch_size, True)
        self.train()
        for _ in range(self.epochs_num):
            for x, y in data_loader:
                self.optimizer.zero_grad()
                pred = self(x)
                loss = self.loss_fn(pred, y)
                loss.backward()
                self.optimizer.step()
        return self

Опишем, что здесь происходит. У нас есть модель нейронной сети Pytorch, которую мы хотим включить в какой-нибудь конвейер bosk и обучить ее на графическом процессоре. Чтобы превратить нашу модель в блок BaseBlock, нам нужно использовать декоратор auto_block.

Декоратор работает с классами, имеющими методы fit и transform. Он использует имена аргументов методов для создания метаинформации о слотах. Каждое имя аргумента создает свой собственный слот. Если в методе fit представлен аргумент, то слот будет помечен как используемый на этапе FIT. То же самое с методом transform и этапом TRANSFORM.

Мы можем использовать auto_block без каких-либо аргументов, и он создаст стандартный ЦПУ блок со стандартным поведением, о котором вы можете узнать из документации. В нашем случае мы хотим, чтобы этот блок был помечен как работающий на графическом процессоре (cpu=False, gpu=True), способный к распараллеливанию (thread_safe=True), и, поскольку мы не реализовывали инициализацию весов нейронной сети в соответствии с заданным случайным начальным числом, мы отметили random_state_field=None, чтобы ничего не делать в методе block.set_random_state.

Теперь давайте создадим конвейер с нашей моделью. Обратите внимание, что нам нужно зарегистрировать наш пользовательский блок в функциональном построителе. Также нужно сказать, что в настоящее время в bosk данные графического процессора хранятся с использованием массивов jax, которые несовместимы с тензорами Pytorch. Вот почему мы не ставили MoveToBlock в конвейер и реализовали всю логику передачи данных на ГПУ внутри модели NeuralNetwork.

dim = 2

b = FunctionalPipelineBuilder()
X, y = b.Input()(), b.TargetInput()()
rf = b.RFC()(X=X, y=y)
rf_roc_auc = b.RocAuc()(gt_y=y, pred_probas=rf)

# регистрация блока
nn_func_block_wrapper = b.new(NeuralNetwork, dim=dim)
# добавление блока в конвейер
nn = nn_func_block_wrapper(x=X, y=y)

nn_roc_auc = b.RocAuc()(gt_y=y, pred_probas=nn)
stack = b.Stack(['rf', 'nn'], axis=1)(rf=rf, nn=nn)
average = b.Average(axis=1)(X=stack)
roc_auc = b.RocAuc()(gt_y=y, pred_probas=average)
pipeline = b.build(
    {'X': X, 'y': y},
    {'probas': average, 'rf roc-auc': rf_roc_auc, 'total roc-auc': roc_auc, 'nn roc-auc': nn_roc_auc}
)
GraphvizPainter(figure_dpi=100).from_pipeline(pipeline).render('pipeline.jpeg')
display(Image('pipeline.jpeg'))

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

x, y = make_moons(200, noise=0.5)
train_x, test_x, train_y, test_y = train_test_split(x, y, test_size=0.5)
train_data = {'X': train_x, 'y': train_y}
test_data = {'X': test_x, 'y': test_y}

fit_exec = GreedyParallelExecutor(pipeline, Stage.FIT, outputs=['rf roc-auc', 'nn roc-auc', 'total roc-auc'])
fit_res = fit_exec(train_data).numpy()
for key, val in fit_res.items():
    print(f'{key}:', round(val, 3))

rf roc-auc: 1.0
total roc-auc: 0.985
nn roc-auc: 0.923

test_exec = GreedyParallelExecutor(pipeline, Stage.TRANSFORM, outputs=['rf roc-auc', 'nn roc-auc', 'total roc-auc'])
test_res = test_exec(test_data).numpy()
for key, val in test_res.items():
    print(f'{key}:', round(val, 3))

rf roc-auc: 0.838
total roc-auc: 0.859
nn roc-auc: 0.891