mirrors/mmsegmentation
1
0
Fork 0
mirror of https://github.com/open-mmlab/mmsegmentation.git synced 2025-06-03 22:03:48 +08:00
Code Issues Projects Releases Wiki Activity

Update train.py (#428)

Browse Source 
* Update train.py

Add user-defined hooks.

* Update train.py

* Update train.py
This commit is contained in:
gszh 2021-10-30 03:13:40 +08:00 committed by GitHub
parent 4b69af7b13
commit d7f82e5dc8
1 changed files with 16 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
mmseg/apis/train.py
View File

@ -5,7 +5,8 @@ import warnings
import numpy as np import numpy as np
import torch import torch
from mmcv.parallel import MMDataParallel, MMDistributedDataParallel from mmcv.parallel import MMDataParallel, MMDistributedDataParallel
from mmcv.runner import build_optimizer, build_runner from mmcv.runner import HOOKS, build_optimizer, build_runner
from mmcv.utils import build_from_cfg
from mmseg.core import DistEvalHook, EvalHook from mmseg.core import DistEvalHook, EvalHook
from mmseg.datasets import build_dataloader, build_dataset from mmseg.datasets import build_dataloader, build_dataset
@ -113,6 +114,20 @@ def train_segmentor(model,
runner.register_hook( runner.register_hook(
eval_hook(val_dataloader, **eval_cfg), priority='LOW') eval_hook(val_dataloader, **eval_cfg), priority='LOW')
# user-defined hooks
if cfg.get('custom_hooks', None):
custom_hooks = cfg.custom_hooks
assert isinstance(custom_hooks, list), \
f'custom_hooks expect list type, but got {type(custom_hooks)}'
for hook_cfg in cfg.custom_hooks:
assert isinstance(hook_cfg, dict), \
'Each item in custom_hooks expects dict type, but got ' \
f'{type(hook_cfg)}'
hook_cfg = hook_cfg.copy()
priority = hook_cfg.pop('priority', 'NORMAL')
hook = build_from_cfg(hook_cfg, HOOKS)
runner.register_hook(hook, priority=priority)
if cfg.resume_from: if cfg.resume_from:
runner.resume(cfg.resume_from) runner.resume(cfg.resume_from)
elif cfg.load_from: elif cfg.load_from: