Почему колбэки CatBoost останавливают обучение после 1-й итерации и как это исправить

Разбираем работу колбэков CatBoost: из-за возврата None в after_iteration обучение прекращается после первой итерации. Правильный код и трекинг прогресса.

CatBoost предоставляет колбэки, позволяющие встроиться в процесс обучения и реагировать после каждого шага бустинга. Частый сценарий — обновление индикатора прогресса в интерфейсе. Но если подключить колбэк без нюансов, обучение может завершиться уже после самой первой итерации — выглядит как сбой, хотя причина куда прозаичнее.

Воспроизведение: обучение останавливается после одной итерации

Ниже фрагмент, который на каждой итерации увеличивает счётчик и передаёт обработчик в fit. Счётчик растёт, но CatBoost тут же останавливается.

class StepHook:
    def after_iteration(self, evt):
        global tick_count
        tick_count += 1

observer = StepHook()
clf.fit(x_train, eval_set=x_valid, callbacks=[observer])
# Либо создать обработчик прямо при вызове
clf.fit(x_train, eval_set=x_valid, callbacks=[StepHook()])

Что на самом деле происходит и почему

Контракт колбэков в CatBoost опирается на возвращаемое значение after_iteration, чтобы решить, продолжать ли обучение. Для продолжения нужно вернуть True, для остановки — False. Если ничего не вернуть, Python вернёт None, а он трактуется как False. Итог — мгновенная остановка после первого вызова.

Это поведение видно в примерах репозитория CatBoost — например, MetricsCheckerCallback и EarlyStopCallback. В первом случае явно возвращают True, чтобы идти дальше; во втором — булево условие, чтобы остановиться на выбранной итерации.

Как исправить

В конце after_iteration возвращайте True, если хотите продолжать обучение.

class StepHook:
    def after_iteration(self, evt):
        global tick_count
        tick_count += 1
        return True

Хранение состояния внутри колбэка

Откажитесь от глобальных переменных — так код проще понимать. Держите счётчик прогресса в самом объекте и обращайтесь к нему после завершения fit.

class StepTracker:
    def __init__(self):
        self.tally = 0

    def after_iteration(self, evt):
        self.tally += 1
        return True

tracker = StepTracker()
clf.fit(x_train, eval_set=x_valid, callbacks=[tracker])
print(tracker.tally)

Чтение текущей итерации из колбэка

Примеры из репозитория также показывают, что объект info, передаваемый в after_iteration, содержит поле iteration. Можно обновлять прогресс прямо из него.

class IterMeter:
    def __init__(self):
        self.tally = 0

    def after_iteration(self, evt):
        self.tally = evt.iteration
        return True

Образцы из примеров CatBoost

Ниже — шаблоны, иллюстрирующие управление потоком через возвращаемое значение. В первом примере проверяется структура метрик и явно возвращается True, чтобы продолжить обучение.

class MetricsGuard:
    def after_iteration(self, info):
        for dataset_name in ['learn', 'validation_0', 'validation_1']:
            assert dataset_name in info.metrics
            for m_name in metric_list:
                assert m_name in info.metrics[dataset_name]
                assert len(info.metrics[dataset_name][m_name]) == info.iteration
        return True

trainer.fit(ds_train, y_train,
            callbacks=[MetricsGuard()],
            eval_set=[val0, val1])

Следующий шаблон использует возвращаемое значение, чтобы остановить обучение на заданном шаге.

class HaltOnIter:
    def __init__(self, stop_at):
        self._stop_at = stop_at

    def after_iteration(self, info):
        return info.iteration != self._stop_at

trainer.fit(ds_train, y_train, callbacks=[
    HaltOnIter(7),
    HaltOnIter(5),
    HaltOnIter(6)
])

Почему это важно

Понимание того, что цикл обучения контролируется возвращаемым значением колбэка, избавляет от недоумения, когда прогон внезапно завершается после одной итерации. Это ещё и даёт тонкий контроль над процессом: можно продолжать обучение, остановиться в нужный момент или реализовать свои правила остановки — всё через один и тот же механизм.

Итоги

Если нужно продолжать обучение из after_iteration, всегда возвращайте True. Храните прогресс внутри объекта колбэка, чтобы прочитать его позже, а для точной индексации итераций используйте info.iteration. С этими небольшими правками обновление прогресс-бара в интерфейсе или внедрение своей логики в цикл обучения CatBoost становится простым и надёжным.