From b39885d953896e0c8c541c7cc006463e3605fb54 Mon Sep 17 00:00:00 2001 From: Ma Zerun Date: Wed, 19 Jan 2022 18:32:55 +0800 Subject: [PATCH] [Feature] Support K-fold cross-validation (#563) * Support to use `indices` to specify which samples to evaluate. * Add KFoldDataset wrapper * Rename 'K' to 'num_splits' accroding to sklearn * Add `kfold-cross-valid.py` * Add unit tests * Add help doc and docstring --- mmcls/datasets/__init__.py | 4 +- mmcls/datasets/base_dataset.py | 5 + mmcls/datasets/builder.py | 10 +- mmcls/datasets/dataset_wrappers.py | 53 +++ mmcls/datasets/multi_label.py | 3 + tests/test_data/test_builder.py | 152 +++++++- .../test_datasets/test_dataset_wrapper.py | 64 +++- tools/kfold-cross-valid.py | 355 ++++++++++++++++++ 8 files changed, 641 insertions(+), 5 deletions(-) create mode 100644 tools/kfold-cross-valid.py diff --git a/mmcls/datasets/__init__.py b/mmcls/datasets/__init__.py index 64fd5ba55..167fef5cf 100644 --- a/mmcls/datasets/__init__.py +++ b/mmcls/datasets/__init__.py @@ -4,7 +4,7 @@ from .builder import (DATASETS, PIPELINES, SAMPLERS, build_dataloader, build_dataset, build_sampler) from .cifar import CIFAR10, CIFAR100 from .dataset_wrappers import (ClassBalancedDataset, ConcatDataset, - RepeatDataset) + KFoldDataset, RepeatDataset) from .imagenet import ImageNet from .imagenet21k import ImageNet21k from .mnist import MNIST, FashionMNIST @@ -17,5 +17,5 @@ __all__ = [ 'VOC', 'MultiLabelDataset', 'build_dataloader', 'build_dataset', 'DistributedSampler', 'ConcatDataset', 'RepeatDataset', 'ClassBalancedDataset', 'DATASETS', 'PIPELINES', 'ImageNet21k', 'SAMPLERS', - 'build_sampler', 'RepeatAugSampler' + 'build_sampler', 'RepeatAugSampler', 'KFoldDataset' ] diff --git a/mmcls/datasets/base_dataset.py b/mmcls/datasets/base_dataset.py index 3c9edf15b..7a2f31092 100644 --- a/mmcls/datasets/base_dataset.py +++ b/mmcls/datasets/base_dataset.py @@ -118,6 +118,7 @@ class BaseDataset(Dataset, metaclass=ABCMeta): results, metric='accuracy', metric_options=None, + indices=None, logger=None): """Evaluate the dataset. @@ -128,6 +129,8 @@ class BaseDataset(Dataset, metaclass=ABCMeta): metric_options (dict, optional): Options for calculating metrics. Allowed keys are 'topk', 'thrs' and 'average_mode'. Defaults to None. + indices (list, optional): The indices of samples corresponding to + the results. Defaults to None. logger (logging.Logger | str, optional): Logger used for printing related information during evaluation. Defaults to None. Returns: @@ -145,6 +148,8 @@ class BaseDataset(Dataset, metaclass=ABCMeta): eval_results = {} results = np.vstack(results) gt_labels = self.get_gt_labels() + if indices is not None: + gt_labels = gt_labels[indices] num_imgs = len(results) assert len(gt_labels) == num_imgs, 'dataset testing results should '\ 'be of the same length as gt_labels.' diff --git a/mmcls/datasets/builder.py b/mmcls/datasets/builder.py index cae66fa99..544f64d7d 100644 --- a/mmcls/datasets/builder.py +++ b/mmcls/datasets/builder.py @@ -1,4 +1,5 @@ # Copyright (c) OpenMMLab. All rights reserved. +import copy import platform import random from functools import partial @@ -25,7 +26,7 @@ SAMPLERS = Registry('sampler') def build_dataset(cfg, default_args=None): from .dataset_wrappers import (ConcatDataset, RepeatDataset, - ClassBalancedDataset) + ClassBalancedDataset, KFoldDataset) if isinstance(cfg, (list, tuple)): dataset = ConcatDataset([build_dataset(c, default_args) for c in cfg]) elif cfg['type'] == 'RepeatDataset': @@ -34,6 +35,13 @@ def build_dataset(cfg, default_args=None): elif cfg['type'] == 'ClassBalancedDataset': dataset = ClassBalancedDataset( build_dataset(cfg['dataset'], default_args), cfg['oversample_thr']) + elif cfg['type'] == 'KFoldDataset': + cp_cfg = copy.deepcopy(cfg) + if cp_cfg.get('test_mode', None) is None: + cp_cfg['test_mode'] = (default_args or {}).pop('test_mode', False) + cp_cfg['dataset'] = build_dataset(cp_cfg['dataset'], default_args) + cp_cfg.pop('type') + dataset = KFoldDataset(**cp_cfg) else: dataset = build_from_cfg(cfg, DATASETS, default_args) diff --git a/mmcls/datasets/dataset_wrappers.py b/mmcls/datasets/dataset_wrappers.py index 68c234e2f..745c8f149 100644 --- a/mmcls/datasets/dataset_wrappers.py +++ b/mmcls/datasets/dataset_wrappers.py @@ -170,3 +170,56 @@ class ClassBalancedDataset(object): def __len__(self): return len(self.repeat_indices) + + +@DATASETS.register_module() +class KFoldDataset: + """A wrapper of dataset for K-Fold cross-validation. + + K-Fold cross-validation divides all the samples in groups of samples, + called folds, of almost equal sizes. And we use k-1 of folds to do training + and use the fold left to do validation. + + Args: + dataset (:obj:`CustomDataset`): The dataset to be divided. + fold (int): The fold used to do validation. Defaults to 0. + num_splits (int): The number of all folds. Defaults to 5. + test_mode (bool): Use the training dataset or validation dataset. + Defaults to False. + seed (int, optional): The seed to shuffle the dataset before splitting. + If None, not shuffle the dataset. Defaults to None. + """ + + def __init__(self, + dataset, + fold=0, + num_splits=5, + test_mode=False, + seed=None): + self.dataset = dataset + self.CLASSES = dataset.CLASSES + self.test_mode = test_mode + self.num_splits = num_splits + + length = len(dataset) + indices = list(range(length)) + if isinstance(seed, int): + rng = np.random.default_rng(seed) + rng.shuffle(indices) + + test_start = length * fold // num_splits + test_end = length * (fold + 1) // num_splits + if test_mode: + self.indices = indices[test_start:test_end] + else: + self.indices = indices[:test_start] + indices[test_end:] + + def __getitem__(self, idx): + return self.dataset[self.indices[idx]] + + def __len__(self): + return len(self.indices) + + def evaluate(self, *args, **kwargs): + kwargs['indices'] = self.indices + return self.dataset.evaluate(*args, **kwargs) diff --git a/mmcls/datasets/multi_label.py b/mmcls/datasets/multi_label.py index 702493e3a..7838ff5ad 100644 --- a/mmcls/datasets/multi_label.py +++ b/mmcls/datasets/multi_label.py @@ -28,6 +28,7 @@ class MultiLabelDataset(BaseDataset): results, metric='mAP', metric_options=None, + indices=None, logger=None, **deprecated_kwargs): """Evaluate the dataset. @@ -62,6 +63,8 @@ class MultiLabelDataset(BaseDataset): eval_results = {} results = np.vstack(results) gt_labels = self.get_gt_labels() + if indices is not None: + gt_labels = gt_labels[indices] num_imgs = len(results) assert len(gt_labels) == num_imgs, 'dataset testing results should '\ 'be of the same length as gt_labels.' diff --git a/tests/test_data/test_builder.py b/tests/test_data/test_builder.py index 534a52e47..44d348900 100644 --- a/tests/test_data/test_builder.py +++ b/tests/test_data/test_builder.py @@ -1,9 +1,14 @@ +import os.path as osp +from copy import deepcopy from unittest.mock import patch import torch from mmcv.utils import digit_version -from mmcls.datasets import build_dataloader +from mmcls.datasets import ImageNet, build_dataloader, build_dataset +from mmcls.datasets.dataset_wrappers import (ClassBalancedDataset, + ConcatDataset, KFoldDataset, + RepeatDataset) class TestDataloaderBuilder(): @@ -119,3 +124,148 @@ class TestDataloaderBuilder(): expect = torch.tensor( [0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4, 5, 5, 5, 6, 6][1::2]) assert all(torch.cat(list(iter(dataloader))) == expect) + + +class TestDatasetBuilder(): + + @classmethod + def setup_class(cls): + data_prefix = osp.join(osp.dirname(__file__), '../data/dataset') + cls.dataset_cfg = dict( + type='ImageNet', + data_prefix=data_prefix, + ann_file=osp.join(data_prefix, 'ann.txt'), + pipeline=[], + test_mode=False, + ) + + def test_normal_dataset(self): + # Test build + dataset = build_dataset(self.dataset_cfg) + assert isinstance(dataset, ImageNet) + assert dataset.test_mode == self.dataset_cfg['test_mode'] + + # Test default_args + dataset = build_dataset(self.dataset_cfg, {'test_mode': True}) + assert dataset.test_mode == self.dataset_cfg['test_mode'] + + cp_cfg = deepcopy(self.dataset_cfg) + cp_cfg.pop('test_mode') + dataset = build_dataset(cp_cfg, {'test_mode': True}) + assert dataset.test_mode + + def test_concat_dataset(self): + # Test build + dataset = build_dataset([self.dataset_cfg, self.dataset_cfg]) + assert isinstance(dataset, ConcatDataset) + assert dataset.datasets[0].test_mode == self.dataset_cfg['test_mode'] + + # Test default_args + dataset = build_dataset([self.dataset_cfg, self.dataset_cfg], + {'test_mode': True}) + assert dataset.datasets[0].test_mode == self.dataset_cfg['test_mode'] + + cp_cfg = deepcopy(self.dataset_cfg) + cp_cfg.pop('test_mode') + dataset = build_dataset([cp_cfg, cp_cfg], {'test_mode': True}) + assert dataset.datasets[0].test_mode + + def test_repeat_dataset(self): + # Test build + dataset = build_dataset( + dict(type='RepeatDataset', dataset=self.dataset_cfg, times=3)) + assert isinstance(dataset, RepeatDataset) + assert dataset.dataset.test_mode == self.dataset_cfg['test_mode'] + + # Test default_args + dataset = build_dataset( + dict(type='RepeatDataset', dataset=self.dataset_cfg, times=3), + {'test_mode': True}) + assert dataset.dataset.test_mode == self.dataset_cfg['test_mode'] + + cp_cfg = deepcopy(self.dataset_cfg) + cp_cfg.pop('test_mode') + dataset = build_dataset( + dict(type='RepeatDataset', dataset=cp_cfg, times=3), + {'test_mode': True}) + assert dataset.dataset.test_mode + + def test_class_balance_dataset(self): + # Test build + dataset = build_dataset( + dict( + type='ClassBalancedDataset', + dataset=self.dataset_cfg, + oversample_thr=1., + )) + assert isinstance(dataset, ClassBalancedDataset) + assert dataset.dataset.test_mode == self.dataset_cfg['test_mode'] + + # Test default_args + dataset = build_dataset( + dict( + type='ClassBalancedDataset', + dataset=self.dataset_cfg, + oversample_thr=1., + ), {'test_mode': True}) + assert dataset.dataset.test_mode == self.dataset_cfg['test_mode'] + + cp_cfg = deepcopy(self.dataset_cfg) + cp_cfg.pop('test_mode') + dataset = build_dataset( + dict( + type='ClassBalancedDataset', + dataset=cp_cfg, + oversample_thr=1., + ), {'test_mode': True}) + assert dataset.dataset.test_mode + + def test_kfold_dataset(self): + # Test build + dataset = build_dataset( + dict( + type='KFoldDataset', + dataset=self.dataset_cfg, + fold=0, + num_splits=5, + test_mode=False, + )) + assert isinstance(dataset, KFoldDataset) + assert not dataset.test_mode + assert dataset.dataset.test_mode == self.dataset_cfg['test_mode'] + + # Test default_args + dataset = build_dataset( + dict( + type='KFoldDataset', + dataset=self.dataset_cfg, + fold=0, + num_splits=5, + test_mode=False, + ), + default_args={ + 'test_mode': True, + 'classes': [1, 2, 3] + }) + assert not dataset.test_mode + assert dataset.dataset.test_mode == self.dataset_cfg['test_mode'] + assert dataset.dataset.CLASSES == [1, 2, 3] + + cp_cfg = deepcopy(self.dataset_cfg) + cp_cfg.pop('test_mode') + dataset = build_dataset( + dict( + type='KFoldDataset', + dataset=self.dataset_cfg, + fold=0, + num_splits=5, + ), + default_args={ + 'test_mode': True, + 'classes': [1, 2, 3] + }) + # The test_mode in default_args will be passed to KFoldDataset + assert dataset.test_mode + assert not dataset.dataset.test_mode + # Other default_args will be passed to child dataset. + assert dataset.dataset.CLASSES == [1, 2, 3] diff --git a/tests/test_data/test_datasets/test_dataset_wrapper.py b/tests/test_data/test_datasets/test_dataset_wrapper.py index 27a18dcae..2798e1fbb 100644 --- a/tests/test_data/test_datasets/test_dataset_wrapper.py +++ b/tests/test_data/test_datasets/test_dataset_wrapper.py @@ -8,7 +8,20 @@ import numpy as np import pytest from mmcls.datasets import (BaseDataset, ClassBalancedDataset, ConcatDataset, - RepeatDataset) + KFoldDataset, RepeatDataset) + + +def mock_evaluate(results, + metric='accuracy', + metric_options=None, + indices=None, + logger=None): + return dict( + results=results, + metric=metric, + metric_options=metric_options, + indices=indices, + logger=logger) @patch.multiple(BaseDataset, __abstractmethods__=set()) @@ -23,6 +36,8 @@ def construct_toy_multi_label_dataset(length): dataset.data_infos = MagicMock() dataset.data_infos.__len__.return_value = length dataset.get_cat_ids = MagicMock(side_effect=lambda idx: cat_ids_list[idx]) + + dataset.evaluate = MagicMock(side_effect=mock_evaluate) return dataset, cat_ids_list @@ -35,6 +50,7 @@ def construct_toy_single_label_dataset(length): dataset.data_infos = MagicMock() dataset.data_infos.__len__.return_value = length dataset.get_cat_ids = MagicMock(side_effect=lambda idx: cat_ids_list[idx]) + dataset.evaluate = MagicMock(side_effect=mock_evaluate) return dataset, cat_ids_list @@ -107,3 +123,49 @@ def test_class_balanced_dataset(construct_dataset): for idx in np.random.randint(0, len(repeat_factor_dataset), 3): assert repeat_factor_dataset[idx] == bisect.bisect_right( repeat_factors_cumsum, idx) + + +@pytest.mark.parametrize('construct_dataset', [ + 'construct_toy_multi_label_dataset', 'construct_toy_single_label_dataset' +]) +def test_kfold_dataset(construct_dataset): + construct_toy_dataset = eval(construct_dataset) + dataset, _ = construct_toy_dataset(10) + + # test without random seed + train_datasets = [ + KFoldDataset(dataset, fold=i, num_splits=3, test_mode=False) + for i in range(5) + ] + test_datasets = [ + KFoldDataset(dataset, fold=i, num_splits=3, test_mode=True) + for i in range(5) + ] + + assert sum([i.indices for i in test_datasets], []) == list(range(10)) + for train_set, test_set in zip(train_datasets, test_datasets): + train_samples = [train_set[i] for i in range(len(train_set))] + test_samples = [test_set[i] for i in range(len(test_set))] + assert set(train_samples + test_samples) == set(range(10)) + + # test with random seed + train_datasets = [ + KFoldDataset(dataset, fold=i, num_splits=3, test_mode=False, seed=1) + for i in range(5) + ] + test_datasets = [ + KFoldDataset(dataset, fold=i, num_splits=3, test_mode=True, seed=1) + for i in range(5) + ] + + assert sum([i.indices for i in test_datasets], []) != list(range(10)) + assert set(sum([i.indices for i in test_datasets], [])) == set(range(10)) + for train_set, test_set in zip(train_datasets, test_datasets): + train_samples = [train_set[i] for i in range(len(train_set))] + test_samples = [test_set[i] for i in range(len(test_set))] + assert set(train_samples + test_samples) == set(range(10)) + + # test evaluate + for test_set in test_datasets: + eval_inputs = test_set.evaluate(None) + assert eval_inputs['indices'] == test_set.indices diff --git a/tools/kfold-cross-valid.py b/tools/kfold-cross-valid.py new file mode 100644 index 000000000..a881316f5 --- /dev/null +++ b/tools/kfold-cross-valid.py @@ -0,0 +1,355 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import argparse +import copy +import os +import os.path as osp +import time +from datetime import datetime +from pathlib import Path + +import mmcv +import torch +from mmcv import Config, DictAction +from mmcv.runner import get_dist_info, init_dist + +from mmcls import __version__ +from mmcls.apis import init_random_seed, set_random_seed, train_model +from mmcls.datasets import build_dataset +from mmcls.models import build_classifier +from mmcls.utils import collect_env, get_root_logger, load_json_log + +TEST_METRICS = ('precision', 'recall', 'f1_score', 'support', 'mAP', 'CP', + 'CR', 'CF1', 'OP', 'OR', 'OF1', 'accuracy') + +prog_description = """K-Fold cross-validation. + +To start a 5-fold cross-validation experiment: + python tools/kfold-cross-valid.py $CONFIG --num-splits 5 + +To resume a 5-fold cross-validation from an interrupted experiment: + python tools/kfold-cross-valid.py $CONFIG --num-splits 5 --resume-from work_dirs/fold2/latest.pth + +To summarize a 5-fold cross-validation: + python tools/kfold-cross-valid.py $CONFIG --num-splits 5 --summary +""" # noqa: E501 + + +def parse_args(): + parser = argparse.ArgumentParser( + formatter_class=argparse.RawDescriptionHelpFormatter, + description=prog_description) + parser.add_argument('config', help='train config file path') + parser.add_argument( + '--num-splits', type=int, help='The number of all folds.') + parser.add_argument( + '--fold', + type=int, + help='The fold used to do validation. ' + 'If specify, only do an experiment of the specified fold.') + parser.add_argument( + '--summary', + action='store_true', + help='Summarize the k-fold cross-validation results.') + parser.add_argument('--work-dir', help='the dir to save logs and models') + parser.add_argument( + '--resume-from', help='the checkpoint file to resume from') + parser.add_argument( + '--no-validate', + action='store_true', + help='whether not to evaluate the checkpoint during training') + group_gpus = parser.add_mutually_exclusive_group() + group_gpus.add_argument('--device', help='device used for training') + group_gpus.add_argument( + '--gpus', + type=int, + help='number of gpus to use ' + '(only applicable to non-distributed training)') + group_gpus.add_argument( + '--gpu-ids', + type=int, + nargs='+', + help='ids of gpus to use ' + '(only applicable to non-distributed training)') + parser.add_argument('--seed', type=int, default=None, help='random seed') + parser.add_argument( + '--deterministic', + action='store_true', + help='whether to set deterministic options for CUDNN backend.') + parser.add_argument( + '--cfg-options', + nargs='+', + action=DictAction, + help='override some settings in the used config, the key-value pair ' + 'in xxx=yyy format will be merged into config file. If the value to ' + 'be overwritten is a list, it should be like key="[a,b]" or key=a,b ' + 'It also allows nested list/tuple values, e.g. key="[(a,b),(c,d)]" ' + 'Note that the quotation marks are necessary and that no white space ' + 'is allowed.') + parser.add_argument( + '--launcher', + choices=['none', 'pytorch', 'slurm', 'mpi'], + default='none', + help='job launcher') + parser.add_argument('--local_rank', type=int, default=0) + args = parser.parse_args() + if 'LOCAL_RANK' not in os.environ: + os.environ['LOCAL_RANK'] = str(args.local_rank) + + return args + + +def copy_config(old_cfg): + """deepcopy a Config object.""" + new_cfg = Config() + _cfg_dict = copy.deepcopy(old_cfg._cfg_dict) + _filename = copy.deepcopy(old_cfg._filename) + _text = copy.deepcopy(old_cfg._text) + super(Config, new_cfg).__setattr__('_cfg_dict', _cfg_dict) + super(Config, new_cfg).__setattr__('_filename', _filename) + super(Config, new_cfg).__setattr__('_text', _text) + return new_cfg + + +def train_single_fold(args, cfg, fold, distributed, seed): + # create the work_dir for the fold + work_dir = osp.join(cfg.work_dir, f'fold{fold}') + cfg.work_dir = work_dir + + # create work_dir + mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir)) + + # wrap the dataset cfg + train_dataset = dict( + type='KFoldDataset', + fold=fold, + dataset=cfg.data.train, + num_splits=args.num_splits, + seed=seed, + ) + val_dataset = dict( + type='KFoldDataset', + fold=fold, + # Use the same dataset with training. + dataset=copy.deepcopy(cfg.data.train), + num_splits=args.num_splits, + seed=seed, + test_mode=True, + ) + val_dataset['dataset']['pipeline'] = cfg.data.val.pipeline + cfg.data.train = train_dataset + cfg.data.val = val_dataset + cfg.data.test = val_dataset + + # dump config + stem, suffix = osp.basename(args.config).rsplit('.', 1) + cfg.dump(osp.join(cfg.work_dir, f'{stem}_fold{fold}.{suffix}')) + # init the logger before other steps + timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime()) + log_file = osp.join(cfg.work_dir, f'{timestamp}.log') + logger = get_root_logger(log_file=log_file, log_level=cfg.log_level) + + # init the meta dict to record some important information such as + # environment info and seed, which will be logged + meta = dict() + # log env info + env_info_dict = collect_env() + env_info = '\n'.join([(f'{k}: {v}') for k, v in env_info_dict.items()]) + dash_line = '-' * 60 + '\n' + logger.info('Environment info:\n' + dash_line + env_info + '\n' + + dash_line) + meta['env_info'] = env_info + + # log some basic info + logger.info(f'Distributed training: {distributed}') + logger.info(f'Config:\n{cfg.pretty_text}') + logger.info( + f'-------- Cross-validation: [{fold+1}/{args.num_splits}] -------- ') + + # set random seeds + # Use different seed in different folds + logger.info(f'Set random seed to {seed + fold}, ' + f'deterministic: {args.deterministic}') + set_random_seed(seed + fold, deterministic=args.deterministic) + cfg.seed = seed + fold + meta['seed'] = seed + fold + + model = build_classifier(cfg.model) + model.init_weights() + + datasets = [build_dataset(cfg.data.train)] + if len(cfg.workflow) == 2: + val_dataset = copy.deepcopy(cfg.data.val) + val_dataset.pipeline = cfg.data.train.pipeline + datasets.append(build_dataset(val_dataset)) + meta.update( + dict( + mmcls_version=__version__, + config=cfg.pretty_text, + CLASSES=datasets[0].CLASSES, + kfold=dict(fold=fold, num_splits=args.num_splits))) + # add an attribute for visualization convenience + train_model( + model, + datasets, + cfg, + distributed=distributed, + validate=(not args.no_validate), + timestamp=timestamp, + device='cpu' if args.device == 'cpu' else 'cuda', + meta=meta) + + +def summary(args, cfg): + summary = dict() + for fold in range(args.num_splits): + work_dir = Path(cfg.work_dir) / f'fold{fold}' + + # Find the latest training log + log_files = list(work_dir.glob('*.log.json')) + if len(log_files) == 0: + continue + log_file = sorted(log_files)[-1] + + date = datetime.fromtimestamp(log_file.lstat().st_mtime) + summary[fold] = {'date': date.strftime('%Y-%m-%d %H:%M:%S')} + + # Find the latest eval log + json_log = load_json_log(log_file) + epochs = sorted(list(json_log.keys())) + eval_log = {} + + def is_metric_key(key): + for metric in TEST_METRICS: + if metric in key: + return True + return False + + for epoch in epochs[::-1]: + if any(is_metric_key(k) for k in json_log[epoch].keys()): + eval_log = json_log[epoch] + break + + summary[fold]['epoch'] = epoch + summary[fold]['metric'] = { + k: v[0] # the value is a list with only one item. + for k, v in eval_log.items() if is_metric_key(k) + } + show_summary(args, summary) + + +def show_summary(args, summary_data): + try: + from rich.console import Console + from rich.table import Table + except ImportError: + raise ImportError('Please run `pip install rich` to install ' + 'package `rich` to draw the table.') + + console = Console() + table = Table(title=f'{args.num_splits}-fold Cross-validation Summary') + table.add_column('Fold') + metrics = summary_data[0]['metric'].keys() + for metric in metrics: + table.add_column(metric) + table.add_column('Epoch') + table.add_column('Date') + + for fold in range(args.num_splits): + row = [f'{fold+1}'] + if fold not in summary_data: + table.add_row(*row) + continue + for metric in metrics: + metric_value = summary_data[fold]['metric'].get(metric, '') + + def format_value(value): + if isinstance(value, float): + return f'{value:.2f}' + if isinstance(value, (list, tuple)): + return str([format_value(i) for i in value]) + else: + return str(value) + + row.append(format_value(metric_value)) + row.append(str(summary_data[fold]['epoch'])) + row.append(summary_data[fold]['date']) + table.add_row(*row) + + console.print(table) + + +def main(): + args = parse_args() + + cfg = Config.fromfile(args.config) + if args.cfg_options is not None: + cfg.merge_from_dict(args.cfg_options) + # set cudnn_benchmark + if cfg.get('cudnn_benchmark', False): + torch.backends.cudnn.benchmark = True + + # work_dir is determined in this priority: CLI > segment in file > filename + if args.work_dir is not None: + # update configs according to CLI args if args.work_dir is not None + cfg.work_dir = args.work_dir + elif cfg.get('work_dir', None) is None: + # use config filename as default work_dir if cfg.work_dir is None + cfg.work_dir = osp.join('./work_dirs', + osp.splitext(osp.basename(args.config))[0]) + + if args.summary: + summary(args, cfg) + return + + # resume from the previous experiment + if args.resume_from is not None: + cfg.resume_from = args.resume_from + resume_kfold = torch.load(cfg.resume_from).get('meta', + {}).get('kfold', None) + if resume_kfold is None: + raise RuntimeError( + 'No "meta" key in checkpoints or no "kfold" in the meta dict. ' + 'Please check if the resume checkpoint from a k-fold ' + 'cross-valid experiment.') + resume_fold = resume_kfold['fold'] + assert args.num_splits == resume_kfold['num_splits'] + else: + resume_fold = 0 + + if args.gpu_ids is not None: + cfg.gpu_ids = args.gpu_ids + else: + cfg.gpu_ids = range(1) if args.gpus is None else range(args.gpus) + + # init distributed env first, since logger depends on the dist info. + if args.launcher == 'none': + distributed = False + else: + distributed = True + init_dist(args.launcher, **cfg.dist_params) + _, world_size = get_dist_info() + cfg.gpu_ids = range(world_size) + + # init a unified random seed + seed = init_random_seed(args.seed) + + # create work_dir + mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir)) + + if args.fold is not None: + folds = [args.fold] + else: + folds = range(resume_fold, args.num_splits) + + for fold in folds: + cfg_ = copy_config(cfg) + if fold != resume_fold: + cfg_.resume_from = None + train_single_fold(args, cfg_, fold, distributed, seed) + + if args.fold is None: + summary(args, cfg) + + +if __name__ == '__main__': + main()