Spaces:
Build error
Build error
File size: 2,184 Bytes
28c256d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 |
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from mmengine.runner import Runner
from mmengine.hooks import Hook
from mmengine.registry import HOOKS, MODELS, RUNNERS
@HOOKS.register_module()
class PrepareTTAHook(Hook):
"""Wraps `runner.model` with subclass of :class:`BaseTTAModel` in
`before_test`.
Note:
This function will only be used with :obj:`MMFullyShardedDataParallel`.
Args:
tta_cfg (dict): Config dictionary of the test time augmentation model.
"""
def __init__(self, tta_cfg: dict):
self.tta_cfg = tta_cfg
def before_test(self, runner: 'Runner') -> None:
"""Wraps `runner.model` with the subclass of :class:`BaseTTAModel`.
Args:
runner (Runner): The runner of the testing process.
"""
self.tta_cfg['module'] = runner.model # type: ignore
model = MODELS.build(self.tta_cfg)
runner.model = model # type: ignore
def build_runner_with_tta(cfg: dict) -> 'Runner':
"""Builds runner with tta (test time augmentation) transformation and
TTAModel.
Note:
This function will only be used with :obj:`MMFullyShardedDataParallel`.
Args:
cfg (dict): cfg with ``tta_pipeline`` and ``tta_model``
Notes:
This is only an experimental feature. We may refactor the code in the
future.
Returns:
Runner: Runner with tta transformation and TTAModel
"""
assert hasattr(
cfg,
'tta_model'), ('please make sure tta_model is defined in your config.')
assert hasattr(cfg, 'tta_pipeline'), (
'please make sure tta_pipeline is defined in your config.')
cfg['test_dataloader']['dataset']['pipeline'] = cfg['tta_pipeline']
if 'runner_type' in cfg:
runner = RUNNERS.build(cfg)
else:
from mmengine.runner import Runner
runner = Runner.from_cfg(cfg)
runner.register_hook(PrepareTTAHook(tta_cfg=cfg['tta_model']))
return runner
|