mmengine/tests/test_hooks/test_early_stopping_hook.py

# Copyright (c) OpenMMLab. All rights reserved.
import logging
import math
import os.path as osp
import tempfile
from unittest.mock import Mock

import torch
import torch.nn as nn
from torch.utils.data import Dataset

from mmengine.evaluator import BaseMetric
from mmengine.hooks import EarlyStoppingHook
from mmengine.logging import MMLogger
from mmengine.model import BaseModel
from mmengine.optim import OptimWrapper
from mmengine.runner import Runner
from mmengine.testing import RunnerTestCase


class ToyModel(BaseModel):

    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(2, 1)

    def forward(self, inputs, data_sample, mode='tensor'):
        labels = torch.stack(data_sample)
        inputs = torch.stack(inputs)
        outputs = self.linear(inputs)
        if mode == 'tensor':
            return outputs
        elif mode == 'loss':
            loss = (labels - outputs).sum()
            outputs = dict(loss=loss)
            return outputs
        else:
            return outputs


class DummyDataset(Dataset):
    METAINFO = dict()  # type: ignore
    data = torch.randn(12, 2)
    label = torch.ones(12)

    @property
    def metainfo(self):
        return self.METAINFO

    def __len__(self):
        return self.data.size(0)

    def __getitem__(self, index):
        return dict(inputs=self.data[index], data_sample=self.label[index])


class DummyMetric(BaseMetric):

    default_prefix: str = 'test'

    def __init__(self, length):
        super().__init__()
        self.length = length
        self.best_idx = length // 2
        self.cur_idx = 0
        self.vals = [90, 91, 92, 88, 89, 90] * 2

    def process(self, *args, **kwargs):
        self.results.append(0)

    def compute_metrics(self, *args, **kwargs):
        acc = self.vals[self.cur_idx]
        self.cur_idx += 1
        return dict(acc=acc)


def get_mock_runner():
    runner = Mock()
    runner.train_loop = Mock()
    runner.train_loop.stop_training = False
    return runner


class TestEarlyStoppingHook(RunnerTestCase):

    def setUp(self):
        self.temp_dir = tempfile.TemporaryDirectory()

    def tearDown(self):
        # `FileHandler` should be closed in Windows, otherwise we cannot
        # delete the temporary directory
        logging.shutdown()
        MMLogger._instance_dict.clear()
        self.temp_dir.cleanup()

    def test_init(self):

        hook = EarlyStoppingHook(monitor='acc')
        self.assertEqual(hook.rule, 'greater')
        self.assertLess(hook.best_score, 0)

        hook = EarlyStoppingHook(monitor='ACC')
        self.assertEqual(hook.rule, 'greater')
        self.assertLess(hook.best_score, 0)

        hook = EarlyStoppingHook(monitor='mAP_50')
        self.assertEqual(hook.rule, 'greater')
        self.assertLess(hook.best_score, 0)

        hook = EarlyStoppingHook(monitor='loss')
        self.assertEqual(hook.rule, 'less')
        self.assertGreater(hook.best_score, 0)

        hook = EarlyStoppingHook(monitor='Loss')
        self.assertEqual(hook.rule, 'less')
        self.assertGreater(hook.best_score, 0)

        hook = EarlyStoppingHook(monitor='ce_loss')
        self.assertEqual(hook.rule, 'less')
        self.assertGreater(hook.best_score, 0)

        with self.assertRaises(ValueError):
            # `rule` should be passed.
            EarlyStoppingHook(monitor='recall')

        with self.assertRaises(ValueError):
            # Invalid `rule`
            EarlyStoppingHook(monitor='accuracy/top1', rule='the world')

    def test_before_run(self):
        runner = Mock()
        runner.train_loop = object()

        # `train_loop` must contain `stop_training` variable.
        with self.assertRaises(AssertionError):
            hook = EarlyStoppingHook(monitor='accuracy/top1', rule='greater')
            hook.before_run(runner)

    def test_after_val_epoch(self):
        runner = get_mock_runner()
        metrics = {'accuracy/top1': 0.5, 'loss': 0.23}
        hook = EarlyStoppingHook(monitor='acc', rule='greater')

        with self.assertWarns(UserWarning):
            # Skip early stopping process since the evaluation results does not
            # include the key 'acc'
            hook.after_val_epoch(runner, metrics)

        # if `monitor` does not match and strict=True, crash the training.
        with self.assertRaises(RuntimeError):
            metrics = {'accuracy/top1': 0.5, 'loss': 0.23}
            hook = EarlyStoppingHook(
                monitor='acc', rule='greater', strict=True)
            hook.after_val_epoch(runner, metrics)

        # Check largest value
        runner = get_mock_runner()
        metrics = [{'accuracy/top1': i / 9.} for i in range(8)]
        hook = EarlyStoppingHook(monitor='accuracy/top1', rule='greater')
        for metric in metrics:
            hook.after_val_epoch(runner, metric)
            if runner.train_loop.stop_training:
                break
        self.assertAlmostEqual(hook.best_score, 7 / 9)

        # Check smallest value
        runner = get_mock_runner()
        metrics = [{'loss': i / 9.} for i in range(8, 0, -1)]
        hook = EarlyStoppingHook(monitor='loss')
        for metric in metrics:
            hook.after_val_epoch(runner, metric)
            if runner.train_loop.stop_training:
                break
        self.assertAlmostEqual(hook.best_score, 1 / 9)

        # Check stop training
        runner = get_mock_runner()
        metrics = [{'accuracy/top1': i} for i in torch.linspace(98, 99, 8)]
        hook = EarlyStoppingHook(
            monitor='accuracy/top1', rule='greater', min_delta=1)
        for metric in metrics:
            hook.after_val_epoch(runner, metric)
            if runner.train_loop.stop_training:
                break
        self.assertTrue(runner.train_loop.stop_training)

        # Check finite
        runner = get_mock_runner()
        metrics = [{'accuracy/top1': math.inf} for i in range(5)]
        hook = EarlyStoppingHook(
            monitor='accuracy/top1', rule='greater', min_delta=1)
        for metric in metrics:
            hook.after_val_epoch(runner, metric)
            if runner.train_loop.stop_training:
                break
        self.assertTrue(runner.train_loop.stop_training)

        # Check patience
        runner = get_mock_runner()
        metrics = [{'accuracy/top1': i} for i in torch.linspace(98, 99, 8)]
        hook = EarlyStoppingHook(
            monitor='accuracy/top1', rule='greater', min_delta=1, patience=10)
        for metric in metrics:
            hook.after_val_epoch(runner, metric)
            if runner.train_loop.stop_training:
                break
        self.assertFalse(runner.train_loop.stop_training)

        # Check stopping_threshold
        runner = get_mock_runner()
        metrics = [{'accuracy/top1': i} for i in torch.linspace(98, 99, 8)]
        hook = EarlyStoppingHook(
            monitor='accuracy/top1',
            rule='greater',
            stopping_threshold=98.5,
            patience=0)
        for metric in metrics:
            hook.after_val_epoch(runner, metric)
            if runner.train_loop.stop_training:
                break
        self.assertAlmostEqual(hook.best_score.item(), 98 + 4 / 7, places=5)

    def test_with_runner(self):
        max_epoch = 10
        work_dir = osp.join(self.temp_dir.name, 'runner_test')
        early_stop_cfg = dict(
            type='EarlyStoppingHook',
            monitor='test/acc',
            rule='greater',
            min_delta=1,
            patience=3,
        )
        runner = Runner(
            model=ToyModel(),
            work_dir=work_dir,
            train_dataloader=dict(
                dataset=DummyDataset(),
                sampler=dict(type='DefaultSampler', shuffle=True),
                batch_size=3,
                num_workers=0),
            val_dataloader=dict(
                dataset=DummyDataset(),
                sampler=dict(type='DefaultSampler', shuffle=False),
                batch_size=3,
                num_workers=0),
            val_evaluator=dict(type=DummyMetric, length=max_epoch),
            optim_wrapper=OptimWrapper(
                torch.optim.Adam(ToyModel().parameters())),
            train_cfg=dict(
                by_epoch=True, max_epochs=max_epoch, val_interval=1),
            val_cfg=dict(),
            custom_hooks=[early_stop_cfg],
            experiment_name='earlystop_test')
        runner.train()
        self.assertEqual(runner.epoch, 6)