mirror of
https://github.com/open-mmlab/mmclassification.git
synced 2025-06-03 21:53:55 +08:00
[Improve] Clean useless code and reduce unit tests memory usage.
This commit is contained in:
mzr1996
parent
cecff79a79
commit
dd660ed99e
@ -1,4 +0,0 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
from .test import ONNXRuntimeClassifier, TensorRTClassifier
|
||||
|
||||
__all__ = ['ONNXRuntimeClassifier', 'TensorRTClassifier']
|
||||
@ -1,96 +0,0 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import warnings
|
||||
|
||||
import numpy as np
|
||||
import onnxruntime as ort
|
||||
import torch
|
||||
|
||||
from mmcls.models.classifiers import BaseClassifier
|
||||
|
||||
|
||||
class ONNXRuntimeClassifier(BaseClassifier):
|
||||
"""Wrapper for classifier's inference with ONNXRuntime."""
|
||||
|
||||
def __init__(self, onnx_file, class_names, device_id):
|
||||
super(ONNXRuntimeClassifier, self).__init__()
|
||||
sess = ort.InferenceSession(onnx_file)
|
||||
|
||||
providers = ['CPUExecutionProvider']
|
||||
options = [{}]
|
||||
is_cuda_available = ort.get_device() == 'GPU'
|
||||
if is_cuda_available:
|
||||
providers.insert(0, 'CUDAExecutionProvider')
|
||||
options.insert(0, {'device_id': device_id})
|
||||
sess.set_providers(providers, options)
|
||||
|
||||
self.sess = sess
|
||||
self.CLASSES = class_names
|
||||
self.device_id = device_id
|
||||
self.io_binding = sess.io_binding()
|
||||
self.output_names = [_.name for _ in sess.get_outputs()]
|
||||
self.is_cuda_available = is_cuda_available
|
||||
|
||||
def simple_test(self, img, img_metas, **kwargs):
|
||||
raise NotImplementedError('This method is not implemented.')
|
||||
|
||||
def extract_feat(self, imgs):
|
||||
raise NotImplementedError('This method is not implemented.')
|
||||
|
||||
def forward_train(self, imgs, **kwargs):
|
||||
raise NotImplementedError('This method is not implemented.')
|
||||
|
||||
def forward_test(self, imgs, img_metas, **kwargs):
|
||||
input_data = imgs
|
||||
# set io binding for inputs/outputs
|
||||
device_type = 'cuda' if self.is_cuda_available else 'cpu'
|
||||
if not self.is_cuda_available:
|
||||
input_data = input_data.cpu()
|
||||
self.io_binding.bind_input(
|
||||
name='input',
|
||||
device_type=device_type,
|
||||
device_id=self.device_id,
|
||||
element_type=np.float32,
|
||||
shape=input_data.shape,
|
||||
buffer_ptr=input_data.data_ptr())
|
||||
|
||||
for name in self.output_names:
|
||||
self.io_binding.bind_output(name)
|
||||
# run session to get outputs
|
||||
self.sess.run_with_iobinding(self.io_binding)
|
||||
results = self.io_binding.copy_outputs_to_cpu()[0]
|
||||
return list(results)
|
||||
|
||||
|
||||
class TensorRTClassifier(BaseClassifier):
|
||||
|
||||
def __init__(self, trt_file, class_names, device_id):
|
||||
super(TensorRTClassifier, self).__init__()
|
||||
from mmcv.tensorrt import TRTWraper, load_tensorrt_plugin
|
||||
try:
|
||||
load_tensorrt_plugin()
|
||||
except (ImportError, ModuleNotFoundError):
|
||||
warnings.warn('If input model has custom op from mmcv, \
|
||||
you may have to build mmcv with TensorRT from source.')
|
||||
model = TRTWraper(
|
||||
trt_file, input_names=['input'], output_names=['probs'])
|
||||
|
||||
self.model = model
|
||||
self.device_id = device_id
|
||||
self.CLASSES = class_names
|
||||
|
||||
def simple_test(self, img, img_metas, **kwargs):
|
||||
raise NotImplementedError('This method is not implemented.')
|
||||
|
||||
def extract_feat(self, imgs):
|
||||
raise NotImplementedError('This method is not implemented.')
|
||||
|
||||
def forward_train(self, imgs, **kwargs):
|
||||
raise NotImplementedError('This method is not implemented.')
|
||||
|
||||
def forward_test(self, imgs, img_metas, **kwargs):
|
||||
input_data = imgs
|
||||
with torch.cuda.device(self.device_id), torch.no_grad():
|
||||
results = self.model({'input': input_data})['probs']
|
||||
results = results.detach().cpu().numpy()
|
||||
|
||||
return list(results)
|
||||
@ -1,12 +1,10 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
from .class_num_check_hook import ClassNumCheckHook
|
||||
from .lr_updater import CosineAnnealingCooldownLrUpdaterHook
|
||||
from .precise_bn_hook import PreciseBNHook
|
||||
from .visualization_hook import VisualizationHook
|
||||
|
||||
__all__ = [
|
||||
'ClassNumCheckHook',
|
||||
'PreciseBNHook',
|
||||
'CosineAnnealingCooldownLrUpdaterHook',
|
||||
'VisualizationHook',
|
||||
]
|
||||
|
||||
@ -1,85 +0,0 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
from math import cos, pi
|
||||
|
||||
from mmcv.runner.hooks import LrUpdaterHook
|
||||
|
||||
from mmcls.registry import HOOKS
|
||||
|
||||
|
||||
@HOOKS.register_module()
|
||||
class CosineAnnealingCooldownLrUpdaterHook(LrUpdaterHook):
|
||||
"""Cosine annealing learning rate scheduler with cooldown.
|
||||
|
||||
Args:
|
||||
min_lr (float, optional): The minimum learning rate after annealing.
|
||||
Defaults to None.
|
||||
min_lr_ratio (float, optional): The minimum learning ratio after
|
||||
nnealing. Defaults to None.
|
||||
cool_down_ratio (float): The cooldown ratio. Defaults to 0.1.
|
||||
cool_down_time (int): The cooldown time. Defaults to 10.
|
||||
by_epoch (bool): If True, the learning rate changes epoch by epoch. If
|
||||
False, the learning rate changes iter by iter. Defaults to True.
|
||||
warmup (string, optional): Type of warmup used. It can be None (use no
|
||||
warmup), 'constant', 'linear' or 'exp'. Defaults to None.
|
||||
warmup_iters (int): The number of iterations or epochs that warmup
|
||||
lasts. Defaults to 0.
|
||||
warmup_ratio (float): LR used at the beginning of warmup equals to
|
||||
``warmup_ratio * initial_lr``. Defaults to 0.1.
|
||||
warmup_by_epoch (bool): If True, the ``warmup_iters``
|
||||
means the number of epochs that warmup lasts, otherwise means the
|
||||
number of iteration that warmup lasts. Defaults to False.
|
||||
|
||||
Note:
|
||||
You need to set one and only one of ``min_lr`` and ``min_lr_ratio``.
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
min_lr=None,
|
||||
min_lr_ratio=None,
|
||||
cool_down_ratio=0.1,
|
||||
cool_down_time=10,
|
||||
**kwargs):
|
||||
assert (min_lr is None) ^ (min_lr_ratio is None)
|
||||
self.min_lr = min_lr
|
||||
self.min_lr_ratio = min_lr_ratio
|
||||
self.cool_down_time = cool_down_time
|
||||
self.cool_down_ratio = cool_down_ratio
|
||||
super(CosineAnnealingCooldownLrUpdaterHook, self).__init__(**kwargs)
|
||||
|
||||
def get_lr(self, runner, base_lr):
|
||||
if self.by_epoch:
|
||||
progress = runner.epoch
|
||||
max_progress = runner.max_epochs
|
||||
else:
|
||||
progress = runner.iter
|
||||
max_progress = runner.max_iters
|
||||
|
||||
if self.min_lr_ratio is not None:
|
||||
target_lr = base_lr * self.min_lr_ratio
|
||||
else:
|
||||
target_lr = self.min_lr
|
||||
|
||||
if progress > max_progress - self.cool_down_time:
|
||||
return target_lr * self.cool_down_ratio
|
||||
else:
|
||||
max_progress = max_progress - self.cool_down_time
|
||||
|
||||
return annealing_cos(base_lr, target_lr, progress / max_progress)
|
||||
|
||||
|
||||
def annealing_cos(start, end, factor, weight=1):
|
||||
"""Calculate annealing cos learning rate.
|
||||
|
||||
Cosine anneal from `weight * start + (1 - weight) * end` to `end` as
|
||||
percentage goes from 0.0 to 1.0.
|
||||
|
||||
Args:
|
||||
start (float): The starting learning rate of the cosine annealing.
|
||||
end (float): The ending learing rate of the cosine annealing.
|
||||
factor (float): The coefficient of `pi` when calculating the current
|
||||
percentage. Range from 0.0 to 1.0.
|
||||
weight (float, optional): The combination factor of `start` and `end`
|
||||
when calculating the actual starting learning rate. Default to 1.
|
||||
"""
|
||||
cos_out = cos(pi * factor) + 1
|
||||
return end + 0.5 * weight * (start - end) * cos_out
|
||||
@ -1,7 +1 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
from .dist_utils import DistOptimizerHook, allreduce_grads, sync_random_seed
|
||||
from .misc import multi_apply
|
||||
|
||||
__all__ = [
|
||||
'allreduce_grads', 'DistOptimizerHook', 'multi_apply', 'sync_random_seed'
|
||||
]
|
||||
|
||||
@ -1,98 +0,0 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
from collections import OrderedDict
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from mmcv.runner import OptimizerHook, get_dist_info
|
||||
from torch._utils import (_flatten_dense_tensors, _take_tensors,
|
||||
_unflatten_dense_tensors)
|
||||
|
||||
|
||||
def _allreduce_coalesced(tensors, world_size, bucket_size_mb=-1):
|
||||
if bucket_size_mb > 0:
|
||||
bucket_size_bytes = bucket_size_mb * 1024 * 1024
|
||||
buckets = _take_tensors(tensors, bucket_size_bytes)
|
||||
else:
|
||||
buckets = OrderedDict()
|
||||
for tensor in tensors:
|
||||
tp = tensor.type()
|
||||
if tp not in buckets:
|
||||
buckets[tp] = []
|
||||
buckets[tp].append(tensor)
|
||||
buckets = buckets.values()
|
||||
|
||||
for bucket in buckets:
|
||||
flat_tensors = _flatten_dense_tensors(bucket)
|
||||
dist.all_reduce(flat_tensors)
|
||||
flat_tensors.div_(world_size)
|
||||
for tensor, synced in zip(
|
||||
bucket, _unflatten_dense_tensors(flat_tensors, bucket)):
|
||||
tensor.copy_(synced)
|
||||
|
||||
|
||||
def allreduce_grads(params, coalesce=True, bucket_size_mb=-1):
|
||||
grads = [
|
||||
param.grad.data for param in params
|
||||
if param.requires_grad and param.grad is not None
|
||||
]
|
||||
world_size = dist.get_world_size()
|
||||
if coalesce:
|
||||
_allreduce_coalesced(grads, world_size, bucket_size_mb)
|
||||
else:
|
||||
for tensor in grads:
|
||||
dist.all_reduce(tensor.div_(world_size))
|
||||
|
||||
|
||||
class DistOptimizerHook(OptimizerHook):
|
||||
|
||||
def __init__(self, grad_clip=None, coalesce=True, bucket_size_mb=-1):
|
||||
self.grad_clip = grad_clip
|
||||
self.coalesce = coalesce
|
||||
self.bucket_size_mb = bucket_size_mb
|
||||
|
||||
def after_train_iter(self, runner):
|
||||
runner.optimizer.zero_grad()
|
||||
runner.outputs['loss'].backward()
|
||||
if self.grad_clip is not None:
|
||||
self.clip_grads(runner.model.parameters())
|
||||
runner.optimizer.step()
|
||||
|
||||
|
||||
def sync_random_seed(seed=None, device='cuda'):
|
||||
"""Make sure different ranks share the same seed.
|
||||
|
||||
All workers must call this function, otherwise it will deadlock.
|
||||
This method is generally used in `DistributedSampler`,
|
||||
because the seed should be identical across all processes
|
||||
in the distributed group.
|
||||
|
||||
In distributed sampling, different ranks should sample non-overlapped
|
||||
data in the dataset. Therefore, this function is used to make sure that
|
||||
each rank shuffles the data indices in the same order based
|
||||
on the same seed. Then different ranks could use different indices
|
||||
to select non-overlapped data from the same data list.
|
||||
|
||||
Args:
|
||||
seed (int, Optional): The seed. Default to None.
|
||||
device (str): The device where the seed will be put on.
|
||||
Default to 'cuda'.
|
||||
|
||||
Returns:
|
||||
int: Seed to be used.
|
||||
"""
|
||||
if seed is None:
|
||||
seed = np.random.randint(2**31)
|
||||
assert isinstance(seed, int)
|
||||
|
||||
rank, world_size = get_dist_info()
|
||||
|
||||
if world_size == 1:
|
||||
return seed
|
||||
|
||||
if rank == 0:
|
||||
random_num = torch.tensor(seed, dtype=torch.int32, device=device)
|
||||
else:
|
||||
random_num = torch.tensor(0, dtype=torch.int32, device=device)
|
||||
dist.broadcast(random_num, src=0)
|
||||
return random_num.item()
|
||||
@ -1,8 +0,0 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
from functools import partial
|
||||
|
||||
|
||||
def multi_apply(func, *args, **kwargs):
|
||||
pfunc = partial(func, **kwargs) if kwargs else func
|
||||
map_results = map(pfunc, *args)
|
||||
return tuple(map(list, zip(*map_results)))
|
||||
@ -1,7 +1,6 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
from .base_dataset import BaseDataset
|
||||
from .builder import (DATASETS, PIPELINES, SAMPLERS, build_dataloader,
|
||||
build_dataset, build_sampler)
|
||||
from .builder import build_dataset
|
||||
from .cifar import CIFAR10, CIFAR100
|
||||
from .cub import CUB
|
||||
from .custom import CustomDataset
|
||||
@ -11,14 +10,12 @@ from .imagenet import ImageNet, ImageNet21k
|
||||
from .mnist import MNIST, FashionMNIST
|
||||
from .multi_label import MultiLabelDataset
|
||||
from .pipelines import * # noqa: F401,F403
|
||||
from .samplers import DistributedSampler, RepeatAugSampler
|
||||
from .samplers import * # noqa: F401,F403
|
||||
from .voc import VOC
|
||||
|
||||
__all__ = [
|
||||
'BaseDataset', 'ImageNet', 'CIFAR10', 'CIFAR100', 'MNIST', 'FashionMNIST',
|
||||
'VOC', 'build_dataloader', 'build_dataset', 'DistributedSampler',
|
||||
'ConcatDataset', 'RepeatDataset', 'ClassBalancedDataset', 'DATASETS',
|
||||
'PIPELINES', 'ImageNet21k', 'SAMPLERS', 'build_sampler',
|
||||
'RepeatAugSampler', 'KFoldDataset', 'CUB', 'CustomDataset',
|
||||
'MultiLabelDataset'
|
||||
'VOC', 'build_dataset', 'ConcatDataset', 'RepeatDataset',
|
||||
'ClassBalancedDataset', 'ImageNet21k', 'KFoldDataset', 'CUB',
|
||||
'CustomDataset', 'MultiLabelDataset'
|
||||
]
|
||||
|
||||
@ -1,184 +1,25 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import copy
|
||||
import platform
|
||||
import random
|
||||
from functools import partial
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from mmcv.parallel import collate
|
||||
from mmcv.runner import get_dist_info
|
||||
from mmcv.utils import digit_version
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
from mmcls.registry import DATA_SAMPLERS, DATASETS, TRANSFORMS
|
||||
|
||||
try:
|
||||
from mmcv.utils import IS_IPU_AVAILABLE
|
||||
except ImportError:
|
||||
IS_IPU_AVAILABLE = False
|
||||
|
||||
if platform.system() != 'Windows':
|
||||
# https://github.com/pytorch/pytorch/issues/973
|
||||
import resource
|
||||
rlimit = resource.getrlimit(resource.RLIMIT_NOFILE)
|
||||
hard_limit = rlimit[1]
|
||||
soft_limit = min(4096, hard_limit)
|
||||
resource.setrlimit(resource.RLIMIT_NOFILE, (soft_limit, hard_limit))
|
||||
|
||||
PIPELINES = TRANSFORMS
|
||||
SAMPLERS = DATA_SAMPLERS
|
||||
from mmcls.registry import DATASETS
|
||||
|
||||
|
||||
def build_dataset(cfg, default_args=None):
|
||||
from .dataset_wrappers import (ClassBalancedDataset, ConcatDataset,
|
||||
KFoldDataset, RepeatDataset)
|
||||
if isinstance(cfg, (list, tuple)):
|
||||
dataset = ConcatDataset([build_dataset(c, default_args) for c in cfg])
|
||||
elif cfg['type'] == 'ConcatDataset':
|
||||
dataset = ConcatDataset(
|
||||
[build_dataset(c, default_args) for c in cfg['datasets']],
|
||||
separate_eval=cfg.get('separate_eval', True))
|
||||
elif cfg['type'] == 'RepeatDataset':
|
||||
dataset = RepeatDataset(
|
||||
build_dataset(cfg['dataset'], default_args), cfg['times'])
|
||||
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 = DATASETS.build(cfg, default_args=default_args)
|
||||
def build_dataset(cfg):
|
||||
"""Build dataset.
|
||||
|
||||
return dataset
|
||||
|
||||
|
||||
def build_dataloader(dataset,
|
||||
samples_per_gpu,
|
||||
workers_per_gpu,
|
||||
num_gpus=1,
|
||||
dist=True,
|
||||
shuffle=True,
|
||||
round_up=True,
|
||||
seed=None,
|
||||
pin_memory=True,
|
||||
persistent_workers=True,
|
||||
sampler_cfg=None,
|
||||
**kwargs):
|
||||
"""Build PyTorch DataLoader.
|
||||
|
||||
In distributed training, each GPU/process has a dataloader.
|
||||
In non-distributed training, there is only one dataloader for all GPUs.
|
||||
|
||||
Args:
|
||||
dataset (Dataset): A PyTorch dataset.
|
||||
samples_per_gpu (int): Number of training samples on each GPU, i.e.,
|
||||
batch size of each GPU.
|
||||
workers_per_gpu (int): How many subprocesses to use for data loading
|
||||
for each GPU.
|
||||
num_gpus (int): Number of GPUs. Only used in non-distributed training.
|
||||
dist (bool): Distributed training/test or not. Default: True.
|
||||
shuffle (bool): Whether to shuffle the data at every epoch.
|
||||
Default: True.
|
||||
round_up (bool): Whether to round up the length of dataset by adding
|
||||
extra samples to make it evenly divisible. Default: True.
|
||||
pin_memory (bool): Whether to use pin_memory in DataLoader.
|
||||
Default: True
|
||||
persistent_workers (bool): If True, the data loader will not shutdown
|
||||
the worker processes after a dataset has been consumed once.
|
||||
This allows to maintain the workers Dataset instances alive.
|
||||
The argument also has effect in PyTorch>=1.7.0.
|
||||
Default: True
|
||||
sampler_cfg (dict): sampler configuration to override the default
|
||||
sampler
|
||||
kwargs: any keyword argument to be used to initialize DataLoader
|
||||
|
||||
Returns:
|
||||
DataLoader: A PyTorch dataloader.
|
||||
Examples:
|
||||
>>> from mmcls.datasets import build_dataset
|
||||
>>> mnist_train = build_dataset(
|
||||
... dict(type='MNIST', data_prefix='data/mnist/', test_mode=False))
|
||||
>>> print(mnist_train)
|
||||
Dataset MNIST
|
||||
Number of samples: 60000
|
||||
Number of categories: 10
|
||||
Prefix of data: data/mnist/
|
||||
>>> mnist_test = build_dataset(
|
||||
... dict(type='MNIST', data_prefix='data/mnist/', test_mode=True))
|
||||
>>> print(mnist_test)
|
||||
Dataset MNIST
|
||||
Number of samples: 10000
|
||||
Number of categories: 10
|
||||
Prefix of data: data/mnist/
|
||||
"""
|
||||
rank, world_size = get_dist_info()
|
||||
|
||||
# Custom sampler logic
|
||||
if sampler_cfg:
|
||||
# shuffle=False when val and test
|
||||
sampler_cfg.update(shuffle=shuffle)
|
||||
sampler = build_sampler(
|
||||
sampler_cfg,
|
||||
default_args=dict(
|
||||
dataset=dataset, num_replicas=world_size, rank=rank,
|
||||
seed=seed))
|
||||
# Default sampler logic
|
||||
elif dist:
|
||||
sampler = build_sampler(
|
||||
dict(
|
||||
type='DistributedSampler',
|
||||
dataset=dataset,
|
||||
num_replicas=world_size,
|
||||
rank=rank,
|
||||
shuffle=shuffle,
|
||||
round_up=round_up,
|
||||
seed=seed))
|
||||
else:
|
||||
sampler = None
|
||||
|
||||
# If sampler exists, turn off dataloader shuffle
|
||||
if sampler is not None:
|
||||
shuffle = False
|
||||
|
||||
if dist:
|
||||
batch_size = samples_per_gpu
|
||||
num_workers = workers_per_gpu
|
||||
else:
|
||||
batch_size = num_gpus * samples_per_gpu
|
||||
num_workers = num_gpus * workers_per_gpu
|
||||
|
||||
init_fn = partial(
|
||||
worker_init_fn, num_workers=num_workers, rank=rank,
|
||||
seed=seed) if seed is not None else None
|
||||
|
||||
if digit_version(torch.__version__) >= digit_version('1.8.0'):
|
||||
kwargs['persistent_workers'] = persistent_workers
|
||||
if IS_IPU_AVAILABLE:
|
||||
from mmcv.device.ipu import IPUDataLoader
|
||||
data_loader = IPUDataLoader(
|
||||
dataset,
|
||||
None,
|
||||
batch_size=samples_per_gpu,
|
||||
num_workers=num_workers,
|
||||
shuffle=shuffle,
|
||||
worker_init_fn=init_fn,
|
||||
**kwargs)
|
||||
else:
|
||||
data_loader = DataLoader(
|
||||
dataset,
|
||||
batch_size=batch_size,
|
||||
sampler=sampler,
|
||||
num_workers=num_workers,
|
||||
collate_fn=partial(collate, samples_per_gpu=samples_per_gpu),
|
||||
pin_memory=pin_memory,
|
||||
shuffle=shuffle,
|
||||
worker_init_fn=init_fn,
|
||||
**kwargs)
|
||||
|
||||
return data_loader
|
||||
|
||||
|
||||
def worker_init_fn(worker_id, num_workers, rank, seed):
|
||||
# The seed of each worker equals to
|
||||
# num_worker * rank + worker_id + user_seed
|
||||
worker_seed = num_workers * rank + worker_id + seed
|
||||
np.random.seed(worker_seed)
|
||||
random.seed(worker_seed)
|
||||
torch.manual_seed(worker_seed)
|
||||
|
||||
|
||||
def build_sampler(cfg, default_args=None):
|
||||
if cfg is None:
|
||||
return None
|
||||
else:
|
||||
return DATA_SAMPLERS.build(cfg, default_args=default_args)
|
||||
return DATASETS.build(cfg)
|
||||
|
||||
@ -1,5 +1,4 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
from .distributed_sampler import DistributedSampler
|
||||
from .repeat_aug import RepeatAugSampler
|
||||
|
||||
__all__ = ('DistributedSampler', 'RepeatAugSampler')
|
||||
__all__ = ('RepeatAugSampler', )
|
||||
|
||||
@ -1,60 +0,0 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import torch
|
||||
from torch.utils.data import DistributedSampler as _DistributedSampler
|
||||
|
||||
from mmcls.core.utils import sync_random_seed
|
||||
from mmcls.registry import DATA_SAMPLERS
|
||||
|
||||
|
||||
@DATA_SAMPLERS.register_module()
|
||||
class DistributedSampler(_DistributedSampler):
|
||||
|
||||
def __init__(self,
|
||||
dataset,
|
||||
num_replicas=None,
|
||||
rank=None,
|
||||
shuffle=True,
|
||||
round_up=True,
|
||||
seed=0):
|
||||
super().__init__(dataset, num_replicas=num_replicas, rank=rank)
|
||||
self.shuffle = shuffle
|
||||
self.round_up = round_up
|
||||
if self.round_up:
|
||||
self.total_size = self.num_samples * self.num_replicas
|
||||
else:
|
||||
self.total_size = len(self.dataset)
|
||||
|
||||
# In distributed sampling, different ranks should sample
|
||||
# non-overlapped data in the dataset. Therefore, this function
|
||||
# is used to make sure that each rank shuffles the data indices
|
||||
# in the same order based on the same seed. Then different ranks
|
||||
# could use different indices to select non-overlapped data from the
|
||||
# same data list.
|
||||
self.seed = sync_random_seed(seed)
|
||||
|
||||
def __iter__(self):
|
||||
# deterministically shuffle based on epoch
|
||||
if self.shuffle:
|
||||
g = torch.Generator()
|
||||
# When :attr:`shuffle=True`, this ensures all replicas
|
||||
# use a different random ordering for each epoch.
|
||||
# Otherwise, the next iteration of this sampler will
|
||||
# yield the same ordering.
|
||||
g.manual_seed(self.epoch + self.seed)
|
||||
indices = torch.randperm(len(self.dataset), generator=g).tolist()
|
||||
else:
|
||||
indices = torch.arange(len(self.dataset)).tolist()
|
||||
|
||||
# add extra samples to make it evenly divisible
|
||||
if self.round_up:
|
||||
indices = (
|
||||
indices *
|
||||
int(self.total_size / len(indices) + 1))[:self.total_size]
|
||||
assert len(indices) == self.total_size
|
||||
|
||||
# subsample
|
||||
indices = indices[self.rank:self.total_size:self.num_replicas]
|
||||
if self.round_up:
|
||||
assert len(indices) == self.num_samples
|
||||
|
||||
return iter(indices)
|
||||
@ -2,9 +2,9 @@ import math
|
||||
|
||||
import torch
|
||||
from mmcv.runner import get_dist_info
|
||||
from mmengine.dist import sync_random_seed
|
||||
from torch.utils.data import Sampler
|
||||
|
||||
from mmcls.core.utils import sync_random_seed
|
||||
from mmcls.registry import DATA_SAMPLERS
|
||||
|
||||
|
||||
|
||||
@ -9,7 +9,7 @@ from unittest.mock import MagicMock, call, patch
|
||||
import numpy as np
|
||||
from mmengine.registry import TRANSFORMS
|
||||
|
||||
from mmcls.datasets import DATASETS
|
||||
from mmcls.registry import DATASETS
|
||||
from mmcls.utils import get_root_logger
|
||||
|
||||
mmcls_logger = get_root_logger()
|
||||
|
||||
@ -31,19 +31,3 @@ def test_convert_to_one_hot():
|
||||
ori_one_hot_targets.scatter_(1, targets.long(), 1)
|
||||
one_hot_targets = convert_to_one_hot(targets, classes)
|
||||
assert torch.equal(ori_one_hot_targets, one_hot_targets)
|
||||
|
||||
|
||||
# test cuda version
|
||||
@pytest.mark.skipif(
|
||||
not torch.cuda.is_available(), reason='requires CUDA support')
|
||||
def test_convert_to_one_hot_cuda():
|
||||
# test with original impl
|
||||
classes = 10
|
||||
targets = torch.randint(high=classes, size=(10, 1)).cuda()
|
||||
ori_one_hot_targets = torch.zeros((targets.shape[0], classes),
|
||||
dtype=torch.long,
|
||||
device=targets.device)
|
||||
ori_one_hot_targets.scatter_(1, targets.long(), 1)
|
||||
one_hot_targets = convert_to_one_hot(targets, classes)
|
||||
assert torch.equal(ori_one_hot_targets, one_hot_targets)
|
||||
assert ori_one_hot_targets.device == one_hot_targets.device
|
||||
|
||||
@ -51,11 +51,11 @@ def test_conformer_backbone():
|
||||
|
||||
assert check_norm_state(model.modules(), True)
|
||||
|
||||
imgs = torch.randn(3, 3, 224, 224)
|
||||
imgs = torch.randn(1, 3, 224, 224)
|
||||
conv_feature, transformer_feature = model(imgs)[-1]
|
||||
assert conv_feature.shape == (3, 64 * 1 * 4
|
||||
assert conv_feature.shape == (1, 64 * 1 * 4
|
||||
) # base_channels * channel_ratio * 4
|
||||
assert transformer_feature.shape == (3, 384)
|
||||
assert transformer_feature.shape == (1, 384)
|
||||
|
||||
# Test Conformer with irregular input size.
|
||||
model = Conformer(**cfg_ori)
|
||||
@ -64,17 +64,17 @@ def test_conformer_backbone():
|
||||
|
||||
assert check_norm_state(model.modules(), True)
|
||||
|
||||
imgs = torch.randn(3, 3, 241, 241)
|
||||
imgs = torch.randn(1, 3, 241, 241)
|
||||
conv_feature, transformer_feature = model(imgs)[-1]
|
||||
assert conv_feature.shape == (3, 64 * 1 * 4
|
||||
assert conv_feature.shape == (1, 64 * 1 * 4
|
||||
) # base_channels * channel_ratio * 4
|
||||
assert transformer_feature.shape == (3, 384)
|
||||
assert transformer_feature.shape == (1, 384)
|
||||
|
||||
imgs = torch.randn(3, 3, 321, 221)
|
||||
imgs = torch.randn(1, 3, 321, 221)
|
||||
conv_feature, transformer_feature = model(imgs)[-1]
|
||||
assert conv_feature.shape == (3, 64 * 1 * 4
|
||||
assert conv_feature.shape == (1, 64 * 1 * 4
|
||||
) # base_channels * channel_ratio * 4
|
||||
assert transformer_feature.shape == (3, 384)
|
||||
assert transformer_feature.shape == (1, 384)
|
||||
|
||||
# Test custom arch Conformer without output cls token
|
||||
cfg = deepcopy(cfg_ori)
|
||||
@ -88,8 +88,8 @@ def test_conformer_backbone():
|
||||
cfg['base_channels'] = 32
|
||||
model = Conformer(**cfg)
|
||||
conv_feature, transformer_feature = model(imgs)[-1]
|
||||
assert conv_feature.shape == (3, 32 * 3 * 4)
|
||||
assert transformer_feature.shape == (3, 128)
|
||||
assert conv_feature.shape == (1, 32 * 3 * 4)
|
||||
assert transformer_feature.shape == (1, 128)
|
||||
|
||||
# Test Conformer with multi out indices
|
||||
cfg = deepcopy(cfg_ori)
|
||||
@ -99,13 +99,13 @@ def test_conformer_backbone():
|
||||
assert len(outs) == 3
|
||||
# stage 1
|
||||
conv_feature, transformer_feature = outs[0]
|
||||
assert conv_feature.shape == (3, 64 * 1)
|
||||
assert transformer_feature.shape == (3, 384)
|
||||
assert conv_feature.shape == (1, 64 * 1)
|
||||
assert transformer_feature.shape == (1, 384)
|
||||
# stage 2
|
||||
conv_feature, transformer_feature = outs[1]
|
||||
assert conv_feature.shape == (3, 64 * 1 * 2)
|
||||
assert transformer_feature.shape == (3, 384)
|
||||
assert conv_feature.shape == (1, 64 * 1 * 2)
|
||||
assert transformer_feature.shape == (1, 384)
|
||||
# stage 3
|
||||
conv_feature, transformer_feature = outs[2]
|
||||
assert conv_feature.shape == (3, 64 * 1 * 4)
|
||||
assert transformer_feature.shape == (3, 384)
|
||||
assert conv_feature.shape == (1, 64 * 1 * 4)
|
||||
assert transformer_feature.shape == (1, 384)
|
||||
|
||||
@ -16,7 +16,7 @@ class TestDeiT(TestCase):
|
||||
|
||||
def setUp(self):
|
||||
self.cfg = dict(
|
||||
arch='deit-base', img_size=224, patch_size=16, drop_rate=0.1)
|
||||
arch='deit-tiny', img_size=224, patch_size=16, drop_rate=0.1)
|
||||
|
||||
def test_init_weights(self):
|
||||
# test weight init cfg
|
||||
@ -60,7 +60,7 @@ class TestDeiT(TestCase):
|
||||
os.remove(checkpoint)
|
||||
|
||||
def test_forward(self):
|
||||
imgs = torch.randn(3, 3, 224, 224)
|
||||
imgs = torch.randn(1, 3, 224, 224)
|
||||
|
||||
# test with_cls_token=False
|
||||
cfg = deepcopy(self.cfg)
|
||||
@ -77,7 +77,7 @@ class TestDeiT(TestCase):
|
||||
self.assertIsInstance(outs, tuple)
|
||||
self.assertEqual(len(outs), 1)
|
||||
patch_token = outs[-1]
|
||||
self.assertEqual(patch_token.shape, (3, 768, 14, 14))
|
||||
self.assertEqual(patch_token.shape, (1, 192, 14, 14))
|
||||
|
||||
# test with output_cls_token
|
||||
cfg = deepcopy(self.cfg)
|
||||
@ -86,9 +86,9 @@ class TestDeiT(TestCase):
|
||||
self.assertIsInstance(outs, tuple)
|
||||
self.assertEqual(len(outs), 1)
|
||||
patch_token, cls_token, dist_token = outs[-1]
|
||||
self.assertEqual(patch_token.shape, (3, 768, 14, 14))
|
||||
self.assertEqual(cls_token.shape, (3, 768))
|
||||
self.assertEqual(dist_token.shape, (3, 768))
|
||||
self.assertEqual(patch_token.shape, (1, 192, 14, 14))
|
||||
self.assertEqual(cls_token.shape, (1, 192))
|
||||
self.assertEqual(dist_token.shape, (1, 192))
|
||||
|
||||
# test without output_cls_token
|
||||
cfg = deepcopy(self.cfg)
|
||||
@ -98,7 +98,7 @@ class TestDeiT(TestCase):
|
||||
self.assertIsInstance(outs, tuple)
|
||||
self.assertEqual(len(outs), 1)
|
||||
patch_token = outs[-1]
|
||||
self.assertEqual(patch_token.shape, (3, 768, 14, 14))
|
||||
self.assertEqual(patch_token.shape, (1, 192, 14, 14))
|
||||
|
||||
# Test forward with multi out indices
|
||||
cfg = deepcopy(self.cfg)
|
||||
@ -109,14 +109,14 @@ class TestDeiT(TestCase):
|
||||
self.assertEqual(len(outs), 3)
|
||||
for out in outs:
|
||||
patch_token, cls_token, dist_token = out
|
||||
self.assertEqual(patch_token.shape, (3, 768, 14, 14))
|
||||
self.assertEqual(cls_token.shape, (3, 768))
|
||||
self.assertEqual(dist_token.shape, (3, 768))
|
||||
self.assertEqual(patch_token.shape, (1, 192, 14, 14))
|
||||
self.assertEqual(cls_token.shape, (1, 192))
|
||||
self.assertEqual(dist_token.shape, (1, 192))
|
||||
|
||||
# Test forward with dynamic input size
|
||||
imgs1 = torch.randn(3, 3, 224, 224)
|
||||
imgs2 = torch.randn(3, 3, 256, 256)
|
||||
imgs3 = torch.randn(3, 3, 256, 309)
|
||||
imgs1 = torch.randn(1, 3, 224, 224)
|
||||
imgs2 = torch.randn(1, 3, 256, 256)
|
||||
imgs3 = torch.randn(1, 3, 256, 309)
|
||||
cfg = deepcopy(self.cfg)
|
||||
model = DistilledVisionTransformer(**cfg)
|
||||
for imgs in [imgs1, imgs2, imgs3]:
|
||||
@ -126,6 +126,6 @@ class TestDeiT(TestCase):
|
||||
patch_token, cls_token, dist_token = outs[-1]
|
||||
expect_feat_shape = (math.ceil(imgs.shape[2] / 16),
|
||||
math.ceil(imgs.shape[3] / 16))
|
||||
self.assertEqual(patch_token.shape, (3, 768, *expect_feat_shape))
|
||||
self.assertEqual(cls_token.shape, (3, 768))
|
||||
self.assertEqual(dist_token.shape, (3, 768))
|
||||
self.assertEqual(patch_token.shape, (1, 192, *expect_feat_shape))
|
||||
self.assertEqual(cls_token.shape, (1, 192))
|
||||
self.assertEqual(dist_token.shape, (1, 192))
|
||||
|
||||
@ -90,7 +90,7 @@ class TestMLPMixer(TestCase):
|
||||
self.assertFalse(torch.allclose(ori_weight, initialized_weight))
|
||||
|
||||
def test_forward(self):
|
||||
imgs = torch.randn(3, 3, 224, 224)
|
||||
imgs = torch.randn(1, 3, 224, 224)
|
||||
|
||||
# test forward with single out indices
|
||||
cfg = deepcopy(self.cfg)
|
||||
@ -99,7 +99,7 @@ class TestMLPMixer(TestCase):
|
||||
self.assertIsInstance(outs, tuple)
|
||||
self.assertEqual(len(outs), 1)
|
||||
feat = outs[-1]
|
||||
self.assertEqual(feat.shape, (3, 768, 196))
|
||||
self.assertEqual(feat.shape, (1, 768, 196))
|
||||
|
||||
# test forward with multi out indices
|
||||
cfg = deepcopy(self.cfg)
|
||||
@ -109,10 +109,10 @@ class TestMLPMixer(TestCase):
|
||||
self.assertIsInstance(outs, tuple)
|
||||
self.assertEqual(len(outs), 3)
|
||||
for feat in outs:
|
||||
self.assertEqual(feat.shape, (3, 768, 196))
|
||||
self.assertEqual(feat.shape, (1, 768, 196))
|
||||
|
||||
# test with invalid input shape
|
||||
imgs2 = torch.randn(3, 3, 256, 256)
|
||||
imgs2 = torch.randn(1, 3, 256, 256)
|
||||
cfg = deepcopy(self.cfg)
|
||||
model = MlpMixer(**cfg)
|
||||
with self.assertRaisesRegex(AssertionError, 'dynamic input shape.'):
|
||||
|
||||
@ -28,7 +28,7 @@ class TestSwinTransformer(TestCase):
|
||||
|
||||
def setUp(self):
|
||||
self.cfg = dict(
|
||||
arch='b', img_size=224, patch_size=4, drop_path_rate=0.1)
|
||||
arch='tiny', img_size=224, patch_size=4, drop_path_rate=0.1)
|
||||
|
||||
def test_arch(self):
|
||||
# Test invalid default arch
|
||||
@ -135,7 +135,7 @@ class TestSwinTransformer(TestCase):
|
||||
self.assertIsInstance(outs, tuple)
|
||||
self.assertEqual(len(outs), 1)
|
||||
feat = outs[-1]
|
||||
self.assertEqual(feat.shape, (1, 1024, 7, 7))
|
||||
self.assertEqual(feat.shape, (1, 768, 7, 7))
|
||||
|
||||
# test with window_size=12
|
||||
cfg = deepcopy(self.cfg)
|
||||
@ -145,7 +145,7 @@ class TestSwinTransformer(TestCase):
|
||||
self.assertIsInstance(outs, tuple)
|
||||
self.assertEqual(len(outs), 1)
|
||||
feat = outs[-1]
|
||||
self.assertEqual(feat.shape, (1, 1024, 12, 12))
|
||||
self.assertEqual(feat.shape, (1, 768, 12, 12))
|
||||
with self.assertRaisesRegex(AssertionError, r'the window size \(12\)'):
|
||||
model(torch.randn(1, 3, 224, 224))
|
||||
|
||||
@ -158,7 +158,7 @@ class TestSwinTransformer(TestCase):
|
||||
self.assertIsInstance(outs, tuple)
|
||||
self.assertEqual(len(outs), 1)
|
||||
feat = outs[-1]
|
||||
self.assertEqual(feat.shape, (1, 1024, 7, 7))
|
||||
self.assertEqual(feat.shape, (1, 768, 7, 7))
|
||||
|
||||
# test multiple output indices
|
||||
cfg = deepcopy(self.cfg)
|
||||
@ -169,7 +169,7 @@ class TestSwinTransformer(TestCase):
|
||||
self.assertEqual(len(outs), 4)
|
||||
for stride, out in zip([2, 4, 8, 8], outs):
|
||||
self.assertEqual(out.shape,
|
||||
(1, 128 * stride, 56 // stride, 56 // stride))
|
||||
(1, 96 * stride, 56 // stride, 56 // stride))
|
||||
|
||||
# test with checkpoint forward
|
||||
cfg = deepcopy(self.cfg)
|
||||
@ -185,7 +185,7 @@ class TestSwinTransformer(TestCase):
|
||||
self.assertIsInstance(outs, tuple)
|
||||
self.assertEqual(len(outs), 1)
|
||||
feat = outs[-1]
|
||||
self.assertEqual(feat.shape, (1, 1024, 7, 7))
|
||||
self.assertEqual(feat.shape, (1, 768, 7, 7))
|
||||
|
||||
# test with dynamic input shape
|
||||
imgs1 = torch.randn(1, 3, 224, 224)
|
||||
@ -200,7 +200,7 @@ class TestSwinTransformer(TestCase):
|
||||
feat = outs[-1]
|
||||
expect_feat_shape = (math.ceil(imgs.shape[2] / 32),
|
||||
math.ceil(imgs.shape[3] / 32))
|
||||
self.assertEqual(feat.shape, (1, 1024, *expect_feat_shape))
|
||||
self.assertEqual(feat.shape, (1, 768, *expect_feat_shape))
|
||||
|
||||
def test_structure(self):
|
||||
# test drop_path_rate decay
|
||||
|
||||
@ -115,7 +115,7 @@ class TestVisionTransformer(TestCase):
|
||||
os.remove(checkpoint)
|
||||
|
||||
def test_forward(self):
|
||||
imgs = torch.randn(3, 3, 224, 224)
|
||||
imgs = torch.randn(1, 3, 224, 224)
|
||||
|
||||
# test with_cls_token=False
|
||||
cfg = deepcopy(self.cfg)
|
||||
@ -132,7 +132,7 @@ class TestVisionTransformer(TestCase):
|
||||
self.assertIsInstance(outs, tuple)
|
||||
self.assertEqual(len(outs), 1)
|
||||
patch_token = outs[-1]
|
||||
self.assertEqual(patch_token.shape, (3, 768, 14, 14))
|
||||
self.assertEqual(patch_token.shape, (1, 768, 14, 14))
|
||||
|
||||
# test with output_cls_token
|
||||
cfg = deepcopy(self.cfg)
|
||||
@ -141,8 +141,8 @@ class TestVisionTransformer(TestCase):
|
||||
self.assertIsInstance(outs, tuple)
|
||||
self.assertEqual(len(outs), 1)
|
||||
patch_token, cls_token = outs[-1]
|
||||
self.assertEqual(patch_token.shape, (3, 768, 14, 14))
|
||||
self.assertEqual(cls_token.shape, (3, 768))
|
||||
self.assertEqual(patch_token.shape, (1, 768, 14, 14))
|
||||
self.assertEqual(cls_token.shape, (1, 768))
|
||||
|
||||
# test without output_cls_token
|
||||
cfg = deepcopy(self.cfg)
|
||||
@ -152,7 +152,7 @@ class TestVisionTransformer(TestCase):
|
||||
self.assertIsInstance(outs, tuple)
|
||||
self.assertEqual(len(outs), 1)
|
||||
patch_token = outs[-1]
|
||||
self.assertEqual(patch_token.shape, (3, 768, 14, 14))
|
||||
self.assertEqual(patch_token.shape, (1, 768, 14, 14))
|
||||
|
||||
# Test forward with multi out indices
|
||||
cfg = deepcopy(self.cfg)
|
||||
@ -163,13 +163,13 @@ class TestVisionTransformer(TestCase):
|
||||
self.assertEqual(len(outs), 3)
|
||||
for out in outs:
|
||||
patch_token, cls_token = out
|
||||
self.assertEqual(patch_token.shape, (3, 768, 14, 14))
|
||||
self.assertEqual(cls_token.shape, (3, 768))
|
||||
self.assertEqual(patch_token.shape, (1, 768, 14, 14))
|
||||
self.assertEqual(cls_token.shape, (1, 768))
|
||||
|
||||
# Test forward with dynamic input size
|
||||
imgs1 = torch.randn(3, 3, 224, 224)
|
||||
imgs2 = torch.randn(3, 3, 256, 256)
|
||||
imgs3 = torch.randn(3, 3, 256, 309)
|
||||
imgs1 = torch.randn(1, 3, 224, 224)
|
||||
imgs2 = torch.randn(1, 3, 256, 256)
|
||||
imgs3 = torch.randn(1, 3, 256, 309)
|
||||
cfg = deepcopy(self.cfg)
|
||||
model = VisionTransformer(**cfg)
|
||||
for imgs in [imgs1, imgs2, imgs3]:
|
||||
@ -179,5 +179,5 @@ class TestVisionTransformer(TestCase):
|
||||
patch_token, cls_token = outs[-1]
|
||||
expect_feat_shape = (math.ceil(imgs.shape[2] / 16),
|
||||
math.ceil(imgs.shape[3] / 16))
|
||||
self.assertEqual(patch_token.shape, (3, 768, *expect_feat_shape))
|
||||
self.assertEqual(cls_token.shape, (3, 768))
|
||||
self.assertEqual(patch_token.shape, (1, 768, *expect_feat_shape))
|
||||
self.assertEqual(cls_token.shape, (1, 768))
|
||||
|
||||
@ -1,204 +0,0 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import logging
|
||||
import tempfile
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import mmcv.runner
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from mmcv.runner import obj_from_dict
|
||||
from mmcv.runner.hooks import DistEvalHook, EvalHook
|
||||
from torch.utils.data import DataLoader, Dataset
|
||||
|
||||
from mmcls.apis import single_gpu_test
|
||||
|
||||
|
||||
class ExampleDataset(Dataset):
|
||||
|
||||
def __getitem__(self, idx):
|
||||
results = dict(img=torch.tensor([1]), img_metas=dict())
|
||||
return results
|
||||
|
||||
def __len__(self):
|
||||
return 1
|
||||
|
||||
|
||||
class ExampleModel(nn.Module):
|
||||
|
||||
def __init__(self):
|
||||
super(ExampleModel, self).__init__()
|
||||
self.test_cfg = None
|
||||
self.conv = nn.Conv2d(3, 3, 3)
|
||||
|
||||
def forward(self, img, img_metas, test_mode=False, **kwargs):
|
||||
return img
|
||||
|
||||
def train_step(self, data_batch, optimizer):
|
||||
loss = self.forward(**data_batch)
|
||||
return dict(loss=loss)
|
||||
|
||||
|
||||
def test_iter_eval_hook():
|
||||
with pytest.raises(TypeError):
|
||||
test_dataset = ExampleModel()
|
||||
data_loader = [
|
||||
DataLoader(
|
||||
test_dataset,
|
||||
batch_size=1,
|
||||
sampler=None,
|
||||
num_worker=0,
|
||||
shuffle=False)
|
||||
]
|
||||
EvalHook(data_loader, by_epoch=False)
|
||||
|
||||
test_dataset = ExampleDataset()
|
||||
test_dataset.evaluate = MagicMock(return_value=dict(test='success'))
|
||||
loader = DataLoader(test_dataset, batch_size=1)
|
||||
model = ExampleModel()
|
||||
data_loader = DataLoader(
|
||||
test_dataset, batch_size=1, sampler=None, num_workers=0, shuffle=False)
|
||||
optim_cfg = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0005)
|
||||
optimizer = obj_from_dict(optim_cfg, torch.optim,
|
||||
dict(params=model.parameters()))
|
||||
|
||||
# test EvalHook
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
eval_hook = EvalHook(data_loader, by_epoch=False)
|
||||
runner = mmcv.runner.IterBasedRunner(
|
||||
model=model,
|
||||
optimizer=optimizer,
|
||||
work_dir=tmpdir,
|
||||
logger=logging.getLogger(),
|
||||
max_iters=1)
|
||||
runner.register_hook(eval_hook)
|
||||
runner.run([loader], [('train', 1)], 1)
|
||||
test_dataset.evaluate.assert_called_with([torch.tensor([1])],
|
||||
logger=runner.logger)
|
||||
|
||||
|
||||
def test_epoch_eval_hook():
|
||||
with pytest.raises(TypeError):
|
||||
test_dataset = ExampleModel()
|
||||
data_loader = [
|
||||
DataLoader(
|
||||
test_dataset,
|
||||
batch_size=1,
|
||||
sampler=None,
|
||||
num_worker=0,
|
||||
shuffle=False)
|
||||
]
|
||||
EvalHook(data_loader, by_epoch=True)
|
||||
|
||||
test_dataset = ExampleDataset()
|
||||
test_dataset.evaluate = MagicMock(return_value=dict(test='success'))
|
||||
loader = DataLoader(test_dataset, batch_size=1)
|
||||
model = ExampleModel()
|
||||
data_loader = DataLoader(
|
||||
test_dataset, batch_size=1, sampler=None, num_workers=0, shuffle=False)
|
||||
optim_cfg = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0005)
|
||||
optimizer = obj_from_dict(optim_cfg, torch.optim,
|
||||
dict(params=model.parameters()))
|
||||
|
||||
# test EvalHook with interval
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
eval_hook = EvalHook(data_loader, by_epoch=True, interval=2)
|
||||
runner = mmcv.runner.EpochBasedRunner(
|
||||
model=model,
|
||||
optimizer=optimizer,
|
||||
work_dir=tmpdir,
|
||||
logger=logging.getLogger(),
|
||||
max_epochs=2)
|
||||
runner.register_hook(eval_hook)
|
||||
runner.run([loader], [('train', 1)])
|
||||
test_dataset.evaluate.assert_called_once_with([torch.tensor([1])],
|
||||
logger=runner.logger)
|
||||
|
||||
|
||||
def multi_gpu_test(model, data_loader, tmpdir=None, gpu_collect=False):
|
||||
results = single_gpu_test(model, data_loader)
|
||||
return results
|
||||
|
||||
|
||||
@patch('mmcls.apis.multi_gpu_test', multi_gpu_test)
|
||||
def test_dist_eval_hook():
|
||||
with pytest.raises(TypeError):
|
||||
test_dataset = ExampleModel()
|
||||
data_loader = [
|
||||
DataLoader(
|
||||
test_dataset,
|
||||
batch_size=1,
|
||||
sampler=None,
|
||||
num_worker=0,
|
||||
shuffle=False)
|
||||
]
|
||||
DistEvalHook(data_loader, by_epoch=False)
|
||||
|
||||
test_dataset = ExampleDataset()
|
||||
test_dataset.evaluate = MagicMock(return_value=dict(test='success'))
|
||||
loader = DataLoader(test_dataset, batch_size=1)
|
||||
model = ExampleModel()
|
||||
data_loader = DataLoader(
|
||||
test_dataset, batch_size=1, sampler=None, num_workers=0, shuffle=False)
|
||||
optim_cfg = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0005)
|
||||
optimizer = obj_from_dict(optim_cfg, torch.optim,
|
||||
dict(params=model.parameters()))
|
||||
|
||||
# test DistEvalHook
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
p = patch('mmcv.engine.multi_gpu_test', multi_gpu_test)
|
||||
p.start()
|
||||
eval_hook = DistEvalHook(data_loader, by_epoch=False)
|
||||
runner = mmcv.runner.IterBasedRunner(
|
||||
model=model,
|
||||
optimizer=optimizer,
|
||||
work_dir=tmpdir,
|
||||
logger=logging.getLogger(),
|
||||
max_iters=1)
|
||||
runner.register_hook(eval_hook)
|
||||
runner.run([loader], [('train', 1)])
|
||||
test_dataset.evaluate.assert_called_with([torch.tensor([1])],
|
||||
logger=runner.logger)
|
||||
p.stop()
|
||||
|
||||
|
||||
@patch('mmcls.apis.multi_gpu_test', multi_gpu_test)
|
||||
def test_dist_eval_hook_epoch():
|
||||
with pytest.raises(TypeError):
|
||||
test_dataset = ExampleModel()
|
||||
data_loader = [
|
||||
DataLoader(
|
||||
test_dataset,
|
||||
batch_size=1,
|
||||
sampler=None,
|
||||
num_worker=0,
|
||||
shuffle=False)
|
||||
]
|
||||
DistEvalHook(data_loader)
|
||||
|
||||
test_dataset = ExampleDataset()
|
||||
test_dataset.evaluate = MagicMock(return_value=dict(test='success'))
|
||||
loader = DataLoader(test_dataset, batch_size=1)
|
||||
model = ExampleModel()
|
||||
data_loader = DataLoader(
|
||||
test_dataset, batch_size=1, sampler=None, num_workers=0, shuffle=False)
|
||||
optim_cfg = dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0005)
|
||||
optimizer = obj_from_dict(optim_cfg, torch.optim,
|
||||
dict(params=model.parameters()))
|
||||
|
||||
# test DistEvalHook
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
p = patch('mmcv.engine.multi_gpu_test', multi_gpu_test)
|
||||
p.start()
|
||||
eval_hook = DistEvalHook(data_loader, by_epoch=True, interval=2)
|
||||
runner = mmcv.runner.EpochBasedRunner(
|
||||
model=model,
|
||||
optimizer=optimizer,
|
||||
work_dir=tmpdir,
|
||||
logger=logging.getLogger(),
|
||||
max_epochs=2)
|
||||
runner.register_hook(eval_hook)
|
||||
runner.run([loader], [('train', 1)])
|
||||
test_dataset.evaluate.assert_called_with([torch.tensor([1])],
|
||||
logger=runner.logger)
|
||||
p.stop()
|
||||
Loading…
x
Reference in New Issue
Block a user