PBT или подбор гиперпараметров с помощью популяции

Представим следующую достаточно банальную ситуацию: вы студент регионального университета в аграрном регионе, который учится на DS/ML и перед вами поставили достаточно банальную задачу — отличать больное растение от здорового.

Казалось бы, все достаточно просто — с помощью CV обрабатываем фотку + нейросеть для классификации. Если задача обработки изображения достаточно тривиальная (есть много готовых решений и тд), то есть она типичная, то всего сложнее для нас будет подобрать гиперпараметры чтобы нейросеть выдавала нужные нам значения.

Существует достаточно большое количество методов для автоматического подбора гиперпараметров, но у меня своя история с эволюционными методами, так что я сегодня выбрал PBT или Population Based Training или подбор на основе популяции

Что это вообще такое и для чего это надо?

Как я уже описал выше перед нами стоит задача подбора гиперпараметров, так что в этом блоке разберем что такое гиперпараметры и чем они отличаются от обычных параметров.
В качестве достаточно наглядного примера возьмем уравнение линейной регрессии

Параметры это то, что модель меняет сама, в процессе обучения, в случае регрессии а и b это параметры.

import numpy as np
from sklearn.linear_model import LinearRegression

X = np.array([[1], [2], [3], [4], [5]])  
y = np.array([2, 4, 5, 4, 5])            

model = LinearRegression()
model.fit(X, y)

print(f"Коэффициент наклона: {model.coef_[0]}")
print(f"Свободный член: {model.intercept_}")
print(f"Предсказание для X=6: {model.predict([[6]])[0]}")

Выше представлен код обычной линейной регрессии, модель сама подберет параметры a и b.
Ну с параметрами разобрались, теперь перейдем к гиперпараметрам.

import numpy as np
import tensorflow as tf

X = np.array([[1], [2], [3], [4], [5]], dtype=np.float32)
y = np.array([[2], [4], [5], [4], [5]], dtype=np.float32)

model = tf.keras.Sequential([
    tf.keras.layers.Dense(1, activation='linear', input_shape=(1,))
])

model.compile(optimizer='sgd', loss='mse')

model.fit(X, y, epochs=500, verbose=0)  # verbose=0 скрывает вывод

weights = model.layers[0].get_weights()
print(f"Коэффициент наклона: {weights[0][0][0]:.4f}")
print(f"Свободный член: {weights[1][0]:.4f}")
print(f"Предсказание для X=6: {model.predict(np.array([[6]]), verbose=0)[0][0]:.4f}")

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

что есть PBT ?

Population Based Training - это асинхронный метод подбора гиперпараметров, работающий по расписанию.
Для понимания метод достаточно не сложный, чуть отойдем от нашей реальной задачи скажем так в абстракцию.

Представим что у нас есть задача научить 10 обезьянок собирать кокосы. Каждой из них мы дали свои параметры по типу "бегай быстро", "смотри вверх" и тд, после чего на 5 минут отправили их в джунгли. По истечению этого времени мы посмотрели кто из обезьян собрал больше всего и их параметры скопировали на тех кто собрал меньше всего и добавили им свои еще какие-то параметры. Таким образом мы каждый раз будем производить селекцию пока не дойдем до нужного нам результата в сборке кокосов.
Ну на обезьянках разобрали, а теперь формализуем метод с помощью математики.

Математика PBT

Небольшие обозначения
- множество индексов популяции;
- пространство параметров модели (веса нейросети);
- пространство гиперпараметров;
- количество поколений;
- пространство признаков;
- множество классов (для многоклассовой классификации)

Состояние популяции в момент времени t

Инициация

(t = 0) , где где ​ - распределения для инициализации.

Оператор обучения

Для задачи классификации функция потерь - кросс-энтропия:

Для бинарной классификации:

Для многоклассовой классификации:

Оператор одного шага градиентного спуска:

, где - скорость обучения (компонента вектора ), а - обучающая выборка

Функция пригодности (fitness) для классификации

- валидационная выборка
- предсказанный класс
- индикаторная функция (1 если условие истинно, 0 иначе)

Для бинарной классификации с порогом 0.5:

Сортировка (ранжирование по качеству)

Для классификации сортируем по убыванию accuracy (больше - лучше):
Определим биекцию - перестановку индексов, такую что:

В случае равенства может использоваться дополнительный критерий (например, порядковый номер).
Обозначим:

- элита
- худшие , где - доля элиты

Элитизм

Для всех :

Элите переходит в следующее поколение без изменений.

Оператор селекции (Exploit)

Для каждого выбираем донора случайным образом из элиты:
Копируем параметры (веса) от донора:

Оператор мутации (Explore)

Гиперпараметры мутируют, чтобы внести разнообразие

Вариант 1: Адаптивное возмущение

Вариант 2: Мультипликативное возмущение

Вариант 3: Перезагрузка

Итоговое правило обновления
Для каждого :

где - независимые реализации шума (может быть аддитивным, мультипликативным или заменой)

Скрытый текст

Для понимания того как это все работает иногда хочется потыкаться в это и посмотреть на числах, так что если вам тоже такое интересно, то прошу

Исходные данные

- веса модели (2 параметра)
- скорость обучения
- размер популяции
- размер элиты
- поколения
Задача - бинарная классификация (0/1)

Поколение 0 (инициализация)

Поколение 1 (t =1)

тренировка Для классификации используем кросс-энтропию:
, где

Пусть градиент для каждой модели:

Тогда:

Оценка качества

Для классификации M - accuracy на валидационной выборке из 5 примеров:
Валидационные данные:

Считаем предсказания: , порог 0.5
Для модели 1 ():

Совпало 3 из 5:

Для модели 2 ():

Совпало 3 из 5:

Для модели 3 ():

Совпало 2 из 5

Сортировка

Для классификации

При равенстве берем по индексу:

Элита:
Худшие:

Элитизм для i = 1

Эксплуатация и мутация для худших

для i = 2:
донор d = 1
Копирование параметров:

Мутация гиперпараметра (мультипликативная):
Пусть

Для i = 3
Донор d = 1

Состояние после поколения 1

Поколение 2 (t=2)

Пусть на этом шаге градиент для всех одинаков (так как веса одинаковые):

Тогда:

Оценка accuracy на тех же валидационных данных

Для модели 1 ():

Совпало 3 из 5:

Для модели 2 ():

Совпало 3 из 5:

Для модели 3 ():

Совпало 3 и 5:

у всех accuracy = 0,6 При равных значениях оставляем порядок без изменений:

Элита:
Худшие:

Элитизм для i = 1

Мутация:
для i = 3
донор d = 1

Мутация

Финальное состояние

С математикой мы разобрались, так что перейдем к реализации метода на коде

P.S. так как задача по растениям была абстрактной, так что данных у меня нет - взял вместо MNIST

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import numpy as np
import matplotlib.pyplot as plt
import os
import time
import random
from collections import defaultdict

# ========== 1. КОНФИГУРАЦИЯ ==========
class Config:
    # Популяция
    population_size = 4  # 4 воркера
    
    # Данные - MNIST!
    dataset = "MNIST"
    batch_size = 64
    
    # Обучение
    num_epochs = 5  # 5 эпох
    ready_interval = 100  # проверка PBT каждые 100 шагов
    
    # Гиперпараметры
    lr_min = 1e-4
    lr_max = 1e-2
    momentum_min = 0.5
    momentum_max = 0.99
    
    # Мутация
    mutate_prob = 0.5
    mutate_factor_low = 0.7
    mutate_factor_high = 1.5
    
    # Технические
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    seed = 42
    checkpoint_dir = "./pbt_mnist_demo"
    
    # Для печати
    verbose = True


config = Config()
os.makedirs(config.checkpoint_dir, exist_ok=True)

# Устанавливаем seed
torch.manual_seed(config.seed)
np.random.seed(config.seed)
random.seed(config.seed)

print("="*60)
print("PBT ДЕМОНСТРАЦИЯ НА MNIST")
print("="*60)
print(f"Устройство: {config.device}")
print(f"Размер популяции: {config.population_size}")
print(f"Эпох: {config.num_epochs}")
print(f"Данные: {config.dataset}")
print("="*60)


# ========== 2. ЗАГРУЗКА MNIST ==========
def load_mnist(config):
    """Загружает MNIST"""
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    
    # Загружаем
    train_dataset = datasets.MNIST(
        './data', train=True, download=True, transform=transform
    )
    test_dataset = datasets.MNIST(
        './data', train=False, download=True, transform=transform
    )
    
    # Берем подмножество для скорости
    train_dataset = torch.utils.data.Subset(train_dataset, range(5000))
    test_dataset = torch.utils.data.Subset(test_dataset, range(1000))
    
    train_loader = DataLoader(
        train_dataset, 
        batch_size=config.batch_size, 
        shuffle=True,
        num_workers=0
    )
    
    test_loader = DataLoader(
        test_dataset, 
        batch_size=config.batch_size, 
        shuffle=False,
        num_workers=0
    )
    
    print(f"MNIST загружен:")
    print(f"  Train: {len(train_dataset)} изображений")
    print(f"  Test: {len(test_dataset)} изображений")
    print(f"  Размер: 28x28, классов: 10")
    
    return train_loader, test_loader


# ========== 3. МОДЕЛЬ ДЛЯ MNIST (ИСПРАВЛЕННАЯ) ==========
class MNISTNet(nn.Module):
    """
    Исправленная модель для MNIST
    Правильно считаем размеры после сверток
    """
    def __init__(self, dropout_rate=0.3):
        super().__init__()
        
        # Сверточная часть
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)  # 28x28 -> 28x28
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)  # 28x28 -> 28x28
        self.pool = nn.MaxPool2d(2, 2)  # 28x28 -> 14x14
        
        # После двух сверток и пулинга: 64 канала * 7 * 7 = 3136
        # Потому что: 28 -> pool -> 14 -> conv -> 14 -> pool -> 7
        self.fc1 = nn.Linear(64 * 7 * 7, 128)
        self.fc2 = nn.Linear(128, 10)
        
        self.dropout1 = nn.Dropout2d(dropout_rate)
        self.dropout2 = nn.Dropout(dropout_rate)
        
    def forward(self, x):
        # Первый блок
        x = self.conv1(x)
        x = nn.functional.relu(x)
        x = self.pool(x)  # 28 -> 14
        
        # Второй блок
        x = self.conv2(x)
        x = nn.functional.relu(x)
        x = self.pool(x)  # 14 -> 7
        
        x = self.dropout1(x)
        
        # Переход к полносвязным слоям
        x = torch.flatten(x, 1)  # [batch, 64*7*7]
        
        # Полносвязные слои
        x = self.fc1(x)
        x = nn.functional.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        
        return x


# ========== 4. ОЦЕНКА МОДЕЛИ ==========
@torch.no_grad()
def evaluate(model, test_loader, device, max_batches=10):
    """Оценка модели"""
    model.eval()
    correct = 0
    total = 0
    
    for i, (data, target) in enumerate(test_loader):
        if i >= max_batches:
            break
        data, target = data.to(device), target.to(device)
        output = model(data)
        pred = output.argmax(dim=1, keepdim=True)
        correct += pred.eq(target.view_as(pred)).sum().item()
        total += target.size(0)
    
    return correct / total if total > 0 else 0.0


# ========== 5. ХРАНИЛИЩЕ ==========
class Store:
    def __init__(self, checkpoint_dir):
        self.checkpoint_dir = checkpoint_dir
        self.stats = defaultdict(list)
        
    def save(self, worker_id, model, optimizer, step, lr, momentum, accuracy):
        path = f"{self.checkpoint_dir}/worker_{worker_id}.pth"
        torch.save({
            'model': model.state_dict(),
            'optimizer': optimizer.state_dict(),
            'step': step,
            'lr': lr,
            'momentum': momentum,
            'accuracy': accuracy
        }, path)
        
        self.stats[worker_id].append({
            'step': step,
            'accuracy': accuracy,
            'lr': lr,
            'momentum': momentum
        })
    
    def load(self, worker_id, device):
        path = f"{self.checkpoint_dir}/worker_{worker_id}.pth"
        if os.path.exists(path):
            return torch.load(path, map_location=device)
        return None
    
    def get_all_accuracies(self):
        accs = {}
        for worker_id in self.stats:
            if self.stats[worker_id]:
                accs[worker_id] = self.stats[worker_id][-1]['accuracy']
        return accs


# ========== 6. EXPLOIT (ОТБОР УСЕЧЕНИЕМ) ==========
def truncation_selection(worker_id, current_acc, all_accs, top_ratio=0.25):
    """Отбор усечением"""
    if len(all_accs) < config.population_size:
        return None
    
    sorted_workers = sorted(all_accs.items(), key=lambda x: x[1])
    n_top = max(1, int(len(sorted_workers) * top_ratio))
    
    top_workers = [w for w, acc in sorted_workers[-n_top:]]
    bottom_workers = [w for w, acc in sorted_workers[:n_top]]
    
    if worker_id in bottom_workers:
        donor_id = random.choice(top_workers)
        return donor_id
    
    return None


# ========== 7. EXPLORE (МУТАЦИЯ) ==========
def mutate(lr, momentum):
    """Мутация гиперпараметров"""
    new_lr = lr
    new_momentum = momentum
    
    if random.random() < config.mutate_prob:
        factor = random.choice([config.mutate_factor_low, config.mutate_factor_high])
        new_lr = lr * factor
        new_lr = max(config.lr_min, min(config.lr_max, new_lr))
    
    if random.random() < config.mutate_prob:
        factor = random.choice([config.mutate_factor_low, config.mutate_factor_high])
        new_momentum = momentum * factor
        new_momentum = max(config.momentum_min, min(config.momentum_max, new_momentum))
    
    return new_lr, new_momentum


# ========== 8. ОБУЧЕНИЕ ВОРКЕРА ==========
def train_worker(worker_id, train_loader, test_loader, store, config):
    """Обучение одного воркера с PBT"""
    device = config.device
    
    # Инициализация гиперпараметров
    lr = 10 ** np.random.uniform(np.log10(config.lr_min), np.log10(config.lr_max))
    momentum = np.random.uniform(config.momentum_min, config.momentum_max)
    
    print(f"\n--- Worker {worker_id} ---")
    print(f"  Начальные: LR={lr:.6f}, Momentum={momentum:.3f}")
    
    # Модель и оптимизатор
    model = MNISTNet().to(device)
    optimizer = optim.SGD(model.parameters(), lr=lr, momentum=momentum)
    criterion = nn.CrossEntropyLoss()
    
    step = 0
    best_acc = 0.0
    
    for epoch in range(config.num_epochs):
        model.train()
        epoch_loss = 0
        
        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = data.to(device), target.to(device)
            
            # Обучение
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
            
            epoch_loss += loss.item()
            step += 1
            
            # PBT - проверка
            if step % config.ready_interval == 0:
                # Оценка
                acc = evaluate(model, test_loader, device)
                
                if acc > best_acc:
                    best_acc = acc
                
                print(f"  Worker {worker_id}, эпоха {epoch+1}, шаг {step}, acc={acc:.3f}")
                
                # Сохраняем
                store.save(worker_id, model, optimizer, step, lr, momentum, acc)
                
                # Получаем все accuracy
                all_accs = store.get_all_accuracies()
                
                # EXPLOIT
                if len(all_accs) >= config.population_size:
                    donor_id = truncation_selection(worker_id, acc, all_accs)
                    
                    if donor_id is not None and donor_id != worker_id:
                        donor_acc = all_accs.get(donor_id, 0)
                        print(f"  >> Worker {worker_id} (acc={acc:.3f}) копирует Worker {donor_id} (acc={donor_acc:.3f})")
                        
                        # Загружаем донора
                        donor_checkpoint = store.load(donor_id, device)
                        if donor_checkpoint:
                            # Копируем веса
                            model.load_state_dict(donor_checkpoint['model'])
                            optimizer.load_state_dict(donor_checkpoint['optimizer'])
                            
                            # EXPLORE - мутация
                            old_lr, old_momentum = lr, momentum
                            lr, momentum = mutate(
                                donor_checkpoint['lr'],
                                donor_checkpoint['momentum']
                            )
                            
                            # Обновляем оптимизатор
                            for param_group in optimizer.param_groups:
                                param_group['lr'] = lr
                                param_group['momentum'] = momentum
                            
                            print(f"     Мутация: LR {old_lr:.6f}->{lr:.6f}, Momentum {old_momentum:.3f}->{momentum:.3f}")
                            
                            # Оценка после копирования
                            new_acc = evaluate(model, test_loader, device)
                            print(f"     Accuracy после копирования: {new_acc:.3f}")
                            
                            # Сохраняем обновленного
                            store.save(worker_id, model, optimizer, step, lr, momentum, new_acc)
        
        # Конец эпохи
        avg_loss = epoch_loss / len(train_loader)
        if config.verbose:
            print(f"  Worker {worker_id}, эпоха {epoch+1} завершена, loss={avg_loss:.4f}")
    
    print(f"--- Worker {worker_id} финал, лучшая acc={best_acc:.3f} ---")
    return {'worker_id': worker_id, 'best_acc': best_acc}


# ========== 9. ВИЗУАЛИЗАЦИЯ ==========
def visualize_results(results, store, config):
    """Визуализация результатов"""
    fig, axes = plt.subplots(2, 2, figsize=(14, 10))
    
    # 1. Динамика accuracy
    ax = axes[0, 0]
    for worker_id in range(config.population_size):
        if worker_id in store.stats:
            steps = [s['step'] for s in store.stats[worker_id]]
            accs = [s['accuracy'] for s in store.stats[worker_id]]
            ax.plot(steps, accs, 'o-', label=f'Worker {worker_id}', linewidth=2, markersize=4)
    
    ax.set_xlabel('Шаг')
    ax.set_ylabel('Accuracy')
    ax.set_title('Динамика accuracy')
    ax.legend()
    ax.grid(True, alpha=0.3)
    ax.set_ylim([0, 1])
    
    # 2. Финальные accuracy
    ax = axes[0, 1]
    workers = []
    final_accs = []
    colors = []
    
    sorted_results = sorted(results, key=lambda x: x['best_acc'], reverse=True)
    
    for i, r in enumerate(sorted_results):
        workers.append(f'W{r["worker_id"]}')
        final_accs.append(r['best_acc'])
        colors.append('green' if i == 0 else 'skyblue')
    
    bars = ax.bar(workers, final_accs, color=colors)
    ax.set_xlabel('Воркер')
    ax.set_ylabel('Финальная accuracy')
    ax.set_title('Итоговое качество')
    ax.set_ylim([0, 1])
    
    for bar, acc in zip(bars, final_accs):
        ax.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.02,
                f'{acc:.3f}', ha='center', fontsize=11)
    
    # 3. Эволюция learning rate
    ax = axes[1, 0]
    for worker_id in range(config.population_size):
        if worker_id in store.stats:
            steps = [s['step'] for s in store.stats[worker_id]]
            lrs = [s['lr'] for s in store.stats[worker_id]]
            ax.plot(steps, lrs, 'o-', label=f'Worker {worker_id}', linewidth=2, markersize=4)
    
    ax.set_xlabel('Шаг')
    ax.set_ylabel('Learning Rate')
    ax.set_title('Эволюция learning rate')
    ax.legend()
    ax.grid(True, alpha=0.3)
    ax.set_yscale('log')
    
    # 4. Эволюция momentum
    ax = axes[1, 1]
    for worker_id in range(config.population_size):
        if worker_id in store.stats:
            steps = [s['step'] for s in store.stats[worker_id]]
            momentums = [s['momentum'] for s in store.stats[worker_id]]
            ax.plot(steps, momentums, 'o-', label=f'Worker {worker_id}', linewidth=2, markersize=4)
    
    ax.set_xlabel('Шаг')
    ax.set_ylabel('Momentum')
    ax.set_title('Эволюция momentum')
    ax.legend()
    ax.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('pbt_mnist_results.png', dpi=150)
    plt.show()
    
    print("\nГрафики сохранены в 'pbt_mnist_results.png'")


# ========== 10. ОСНОВНАЯ ФУНКЦИЯ ==========
def run_pbt_demo():
    """Запуск демонстрации PBT"""
    print("\n[1] Загрузка данных MNIST...")
    train_loader, test_loader = load_mnist(config)
    
    print("\n[2] Инициализация хранилища...")
    store = Store(config.checkpoint_dir)
    
    print("\n[3] Запуск популяции воркеров...")
    results = []
    start_time = time.time()
    
    for worker_id in range(config.population_size):
        result = train_worker(worker_id, train_loader, test_loader, store, config)
        results.append(result)
    
    elapsed = time.time() - start_time
    
    print("\n" + "="*60)
    print("РЕЗУЛЬТАТЫ")
    print("="*60)
    
    sorted_results = sorted(results, key=lambda x: x['best_acc'], reverse=True)
    
    for i, r in enumerate(sorted_results):
        print(f"{i+1}. Worker {r['worker_id']}: accuracy = {r['best_acc']:.3f}")
    
    best = sorted_results[0]
    print("-"*40)
    print(f"ПОБЕДИТЕЛЬ: Worker {best['worker_id']} с accuracy {best['best_acc']:.3f}")
    print(f"Время выполнения: {elapsed:.1f} секунд ({elapsed/60:.1f} минут)")
    
    print("\n[4] Визуализация...")
    visualize_results(results, store, config)
    
    print("\n" + "="*60)
    print("ДЕМОНСТРАЦИЯ PBT ЗАВЕРШЕНА")
    print("="*60)


# ========== ЗАПУСК ==========
if __name__ == "__main__":
    run_pbt_demo()

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