Примеры использования Adept

Обучение MLP на MNIST

Полный пример обучения многослойного перцептрона (MLP) для классификации изображений из датасета MNIST с поддержкой CPU/GPU, сохранением чекпоинтов и тестированием.

1. Импорт зависимостей

import sys
from pathlib import Path
from tqdm import tqdm
import time
from argparse import ArgumentParser

from adept import (
    Variable,
    Tensor,
    TensorProperties,
    Shape,
    device_t,
    dtype_t,
    no_grad,
)
from adept.nn import Linear, Module, relu
from adept.optim import SGD
from adept.data import MNISTDataset, CPPDataLoader
from adept.loss import cross_entropy_with_logits
from adept.serialize import FileInput, FileOutput

Что импортируем:

Variable — тензор с поддержкой автоматического дифференцирования
Tensor, TensorProperties, Shape — базовые структуры для работы с тензорами
device_t, dtype_t — типы устройства (CPU/GPU) и данных
no_grad — контекст для отключения градиентов
Linear, Module, relu — модули нейронной сети
SGD — оптимизатор стохастического градиентного спуска
MNISTDataset, CPPDataLoader — загрузчики данных
cross_entropy_with_logits — функция потерь
FileInput, FileOutput — сериализация для чекпоинтов

2. Глобальные параметры

device = device_t.CPU
dtype = dtype_t.Float32

epochs = 5
batch_size = 32
lr = 0.1
lr_decay = 0.5

Настройки:

Устройство по умолчанию — CPU (можно изменить через -g)
Тип данных — Float32
5 эпох обучения
Размер батча — 32
Начальный learning rate — 0.1
Затухание LR — 0.5 после каждой эпохи

3. Определение модели

class MLP(Module):
    def __init__(self):
        super().__init__()
        self.l1 = Linear(28 * 28, 512, dtype=dtype)
        self.l2 = Linear(512, 256, dtype=dtype)
        self.l3 = Linear(256, 10, dtype=dtype)

    def forward(self, x):
        out = relu(self.l1(x))
        out = relu(self.l2(out))
        out = self.l3(out)
        return out

Архитектура:

Вход: 784 нейрона (28×28 пикселей)
Скрытый слой 1: 512 нейронов + ReLU
Скрытый слой 2: 256 нейронов + ReLU
Выход: 10 нейронов (классы цифр 0-9)

4. Парсинг аргументов

def main():
    parser = ArgumentParser()
    parser.add_argument("mnist_path")
    parser.add_argument("-c", "--checkpoint", type=str)
    parser.add_argument("-g", "--gpu", action="store_true")
    args = parser.parse_args()

    mnist_path = Path(args.mnist_path)

    if args.gpu:
        global device
        device = device_t.GPU

Аргументы командной строки:

mnist_path — путь к датасету MNIST
-c, --checkpoint — путь к чекпоинту для продолжения обучения
-g, --gpu — использование GPU вместо CPU

5. Загрузка данных

    train_images_file = str(mnist_path / "train-images.idx3-ubyte")
    train_labels_file = str(mnist_path / "train-labels.idx1-ubyte")

    train_dataset = MNISTDataset(train_images_file, train_labels_file, dtype, True)
    train_dataloader = CPPDataLoader(train_dataset, batch_size)

    test_images_file = str(mnist_path / "t10k-images.idx3-ubyte")
    test_labels_file = str(mnist_path / "t10k-labels.idx1-ubyte")

    test_dataset = MNISTDataset(test_images_file, test_labels_file, dtype, True)
    test_dataloader = CPPDataLoader(test_dataset, batch_size)

Датасеты:

Отдельные загрузчики для тренировочных и тестовых данных
CPPDataLoader — высокопроизводительный загрузчик на C++

6. Инициализация модели и оптимизатора

    mlp = MLP()

    global lr
    if args.checkpoint:
        input = FileInput(args.checkpoint)
        mlp.load(input)
        lr = input.read("lr", dtype_t.Float32)

    mlp.to(device)

    optimizer = SGD(mlp.parameters(), lr)

Что происходит:

Создание экземпляра модели
Загрузка чекпоинта (если указан)
Перенос модели на целевое устройство
Инициализация оптимизатора SGD

7. Цикл обучения

    for epoch in tqdm(range(epochs), unit="epoch"):
        mlp.train()
        pbar = tqdm(train_dataloader, unit="batch")
        for b_i, batch in enumerate(pbar):
            x, y = batch
            out = mlp.forward(Variable(x.to(device), requires_grad=False))

            loss = cross_entropy_with_logits(
                out, Variable(y.to(device), requires_grad=False)
            )
            if b_i % 64 == 0:
                pbar.set_postfix(loss=loss.data().cpu().at([0, 0]))

            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
        optimizer.set_lr(optimizer.lr() * lr_decay)

Шаги обучения:

Переключение модели в режим обучения (train())
Прямой проход (forward pass)
Вычисление функции потерь
Обратное распространение (backward())
Шаг оптимизатора (step())
Обнуление градиентов (zero_grad())
Затухание learning rate после эпохи

8. Сохранение чекпоинта

        # save checkpoint
        checkpoint_path = f"checkpoint_{epoch}_{int(time.time()*1000.0)}.pt"
        output = FileOutput(checkpoint_path)
        mlp.save(output)
        output.write("lr", optimizer.lr())

Сохраняется:

Веса модели
Текущий learning rate
Уникальное имя файла с меткой времени

9. Тестирование

        # test
        mlp.eval()
        with no_grad():
            total_loss = Tensor.zero(TensorProperties(Shape([1]), device, dtype))
            for b_i, batch in enumerate(test_dataloader):
                x, y = batch
                out = mlp.forward(Variable(x.to(device)))
                loss = cross_entropy_with_logits(out, Variable(y.to(device)))
                total_loss += loss.data()
            pbar.set_postfix(test_loss=total_loss.cpu().at([0]) / b_i)

Особенности:

Переключение в режим оценки (eval())
Контекст no_grad() для экономии памяти
Подсчёт среднего значения функции потерь на тестовом наборе

10. Запуск

if __name__ == "__main__":
    main()

Запуск примера

# Обучение на CPU
python mlp_mnist.py /path/to/mnist

# Обучение на GPU
python mlp_mnist.py /path/to/mnist -g

# Продолжение с чекпоинта
python mlp_mnist.py /path/to/mnist -c checkpoint_2_1234567890.pt