From ac99def9426b586e4f6715ddf63db73b04ad8c4e Mon Sep 17 00:00:00 2001 From: Songki Choi Date: Thu, 17 Mar 2022 23:41:12 +0900 Subject: [PATCH 1/2] Remove HPO/MDA/SelfSL features from external release --- init_venv.sh | 2 - mpa/__init__.py | 3 - mpa/modules/datasets/__init__.py | 1 - mpa/modules/datasets/pipelines/__init__.py | 1 - .../datasets/pipelines/selfsl_pipelines.py | 583 ------------------ mpa/modules/datasets/selfsl_dataset.py | 73 --- mpa/modules/models/__init__.py | 4 +- mpa/modules/models/heads/__init__.py | 1 - mpa/modules/models/heads/selfsl_head.py | 96 --- mpa/modules/models/necks/__init__.py | 1 - mpa/modules/models/necks/selfsl_neck.py | 118 ---- mpa/modules/models/selfsl_model.py | 187 ------ mpa/selfsl/__init__.py | 2 - mpa/selfsl/builder.py | 142 ----- mpa/selfsl/trainer.py | 391 ------------ mpa/utils/hpo_stage.py | 240 ------- mpa/utils/hpo_trainer.py | 8 - mpa/utils/mda_stage.py | 178 ------ 18 files changed, 1 insertion(+), 2030 deletions(-) delete mode 100644 mpa/modules/datasets/pipelines/selfsl_pipelines.py delete mode 100644 mpa/modules/datasets/selfsl_dataset.py delete mode 100644 mpa/modules/models/heads/selfsl_head.py delete mode 100644 mpa/modules/models/necks/selfsl_neck.py delete mode 100755 mpa/modules/models/selfsl_model.py delete mode 100755 mpa/selfsl/__init__.py delete mode 100644 mpa/selfsl/builder.py delete mode 100644 mpa/selfsl/trainer.py delete mode 100644 mpa/utils/hpo_stage.py delete mode 100644 mpa/utils/hpo_trainer.py delete mode 100644 mpa/utils/mda_stage.py diff --git a/init_venv.sh b/init_venv.sh index c99e2cbf..917ce4ab 100755 --- a/init_venv.sh +++ b/init_venv.sh @@ -141,8 +141,6 @@ else echo "OTE_PATH should be specified to install dependencies" exit 1 fi -pip install -e external/mda -c ${CONSTRAINTS_FILE} || exit 1 -pip install -e external/hpo -c ${CONSTRAINTS_FILE} || exit 1 deactivate diff --git a/mpa/__init__.py b/mpa/__init__.py index 53bb1143..5012a37c 100644 --- a/mpa/__init__.py +++ b/mpa/__init__.py @@ -7,10 +7,7 @@ from . import cls from . import det from . import seg -from . import selfsl -from .utils.hpo_stage import HpoRunner -from .utils.mda_stage import MdaRunner from .version import __version__, get_version from .builder import build, build_workflow_hook from .stage import Stage, get_available_types diff --git a/mpa/modules/datasets/__init__.py b/mpa/modules/datasets/__init__.py index bfe3a6c0..da504173 100644 --- a/mpa/modules/datasets/__init__.py +++ b/mpa/modules/datasets/__init__.py @@ -5,7 +5,6 @@ from . import det_csv_dataset from . import det_incr_dataset from . import multi_cls_dataset -from . import selfsl_dataset from . import task_adapt_dataset from . import task_adapt_dataset_seg from . import torchvision_dataset diff --git a/mpa/modules/datasets/pipelines/__init__.py b/mpa/modules/datasets/pipelines/__init__.py index d27205a9..dcea2df2 100644 --- a/mpa/modules/datasets/pipelines/__init__.py +++ b/mpa/modules/datasets/pipelines/__init__.py @@ -1,4 +1,3 @@ # flake8: noqa from . import transforms -from . import selfsl_pipelines from . import torchvision2mmdet diff --git a/mpa/modules/datasets/pipelines/selfsl_pipelines.py b/mpa/modules/datasets/pipelines/selfsl_pipelines.py deleted file mode 100644 index 375c0a89..00000000 --- a/mpa/modules/datasets/pipelines/selfsl_pipelines.py +++ /dev/null @@ -1,583 +0,0 @@ -from collections.abc import Sequence -import math -import numbers -import random -import warnings -import numpy as np -import torch -from torchvision.transforms import functional as F -from torchvision import transforms as T -from PIL import Image, ImageFilter -try: - import accimage -except ImportError: - accimage = None - -import mmcv -from mmcv.parallel import DataContainer as DC - -from mpa.selfsl.builder import PIPELINES, build_pipeline - - -_pil_interpolation_to_str = { - Image.NEAREST: 'PIL.Image.NEAREST', - Image.BILINEAR: 'PIL.Image.BILINEAR', - Image.BICUBIC: 'PIL.Image.BICUBIC', - Image.LANCZOS: 'PIL.Image.LANCZOS', - Image.HAMMING: 'PIL.Image.HAMMING', - Image.BOX: 'PIL.Image.BOX', -} - - -def _get_image_size(img): - if F._is_pil_image(img): - return img.size - elif isinstance(img, torch.Tensor) and img.dim() > 2: - return img.shape[-2:][::-1] - else: - raise TypeError("Unexpected type {}".format(type(img))) - - -def to_tensor(data): - """Convert objects of various python types to :obj:`torch.Tensor`. - - Supported types are: :class:`numpy.ndarray`, :class:`torch.Tensor`, - :class:`Sequence`, :class:`int` and :class:`float`. - """ - if isinstance(data, torch.Tensor): - return data - elif isinstance(data, np.ndarray): - return torch.from_numpy(data) - elif isinstance(data, Sequence) and not mmcv.is_str(data): - return torch.tensor(data) - elif isinstance(data, int): - return torch.LongTensor([data]) - elif isinstance(data, float): - return torch.FloatTensor([data]) - else: - raise TypeError( - f'Type {type(data)} cannot be converted to tensor.' - 'Supported types are: `numpy.ndarray`, `torch.Tensor`, ' - '`Sequence`, `int` and `float`') - - -class Lambda: - """Apply a user-defined lambda as a transform. This transform does not support torchscript. - - Args: - lambd (function): Lambda/function to be used for transform. - """ - - def __init__(self, lambd): - if not callable(lambd): - raise TypeError("Argument lambd should be callable, got {}".format(repr(type(lambd).__name__))) - self.lambd = lambd - - def __call__(self, img): - return self.lambd(img) - - def __repr__(self): - return self.__class__.__name__ + '()' - - -@PIPELINES.register_module() -class Compose(object): - """Compose a data pipeline with a sequence of transforms. - - Args: - transforms (list[dict | callable]): - Either config dicts of transforms or transform objects. - """ - - def __init__(self, transforms): - if not isinstance(transforms, Sequence): - raise TypeError(f'transforms must be a type of Sequence, but got {type(transforms)}') - self.transforms = [] - for transform in transforms: - if isinstance(transform, dict): - transform = build_pipeline(transform) - self.transforms.append(transform) - elif callable(transform): - self.transforms.append(transform) - else: - raise TypeError('transform must be callable or a dict, but got' - f' {type(transform)}') - - def __call__(self, data): - for t in self.transforms: - data = t(data) - if data is None: - return None - return data - - def __repr__(self): - format_string = self.__class__.__name__ + '(' - for t in self.transforms: - format_string += f'\n {t}' - format_string += '\n)' - return format_string - - -@PIPELINES.register_module() -class RandomResizedCrop(object): - """Crop the given PIL Image to random size and aspect ratio. - A crop of random size (default: of 0.08 to 1.0) of the original size and a random - aspect ratio (default: of 3/4 to 4/3) of the original aspect ratio is made. This crop - is finally resized to given size. - This is popularly used to train the Inception networks. - Args: - size: expected output size of each edge - scale: range of size of the origin size cropped - ratio: range of aspect ratio of the origin aspect ratio cropped - interpolation: Default: PIL.Image.BILINEAR - """ - - def __init__(self, size, scale=(0.08, 1.0), ratio=(3./4., 4./3.), interpolation=Image.BILINEAR, with_coord=False): - if isinstance(size, (tuple, list)): - self.size = size - else: - self.size = (size, size) - if (scale[0] > scale[1]) or (ratio[0] > ratio[1]): - warnings.warn("range should be of kind (min, max)") - - self.interpolation = interpolation - self.scale = scale - self.ratio = ratio - - self.with_coord = with_coord - - @staticmethod - def get_params(img, scale, ratio): - """Get parameters for ``crop`` for a random sized crop. - Args: - img (PIL Image): Image to be cropped. - scale (tuple): range of size of the origin size cropped - ratio (tuple): range of aspect ratio of the origin aspect ratio cropped - Returns: - tuple: params (i, j, h, w) to be passed to ``crop`` for a random - sized crop. - """ - width, height = _get_image_size(img) - area = height * width - - for attempt in range(10): - target_area = random.uniform(*scale) * area - log_ratio = (math.log(ratio[0]), math.log(ratio[1])) - aspect_ratio = math.exp(random.uniform(*log_ratio)) - - w = int(round(math.sqrt(target_area * aspect_ratio))) - h = int(round(math.sqrt(target_area / aspect_ratio))) - - if 0 < w <= width and 0 < h <= height: - i = random.randint(0, height - h) - j = random.randint(0, width - w) - return i, j, h, w, height, width - - # Fallback to central crop - in_ratio = float(width) / float(height) - if (in_ratio < min(ratio)): - w = width - h = int(round(w / min(ratio))) - elif (in_ratio > max(ratio)): - h = height - w = int(round(h * max(ratio))) - else: # whole image - w = width - h = height - i = (height - h) // 2 - j = (width - w) // 2 - return i, j, h, w, height, width - - def __call__(self, results): - """ - Args: - img (PIL Image): Image to be cropped and resized. - Returns: - PIL Image: Randomly cropped and resized image. - """ - for key in results.get('img_fields', ['img']): - img = results[key] - i, j, h, w, height, width = self.get_params(img, self.scale, self.ratio) - if self.with_coord: - results['coord'] = torch.Tensor([float(j) / (width - 1), float(i) / (height - 1), - float(j + w - 1) / (width - 1), float(i + h - 1) / (height - 1)]) - results[key] = F.resized_crop(img, i, j, h, w, self.size, self.interpolation) - return results - - def __repr__(self): - interpolate_str = _pil_interpolation_to_str[self.interpolation] - format_string = self.__class__.__name__ + '(size={0}'.format(self.size) - format_string += ', scale={0}'.format(tuple(round(s, 4) for s in self.scale)) - format_string += ', ratio={0}'.format(tuple(round(r, 4) for r in self.ratio)) - format_string += ', interpolation={0})'.format(interpolate_str) - format_string += ', with_coord={0}'.format(self.with_coord) - return format_string - - -@PIPELINES.register_module() -class RandomHorizontalFlip(object): - """Horizontally flip the given PIL Image randomly with a given probability. - Args: - p (float): probability of the image being flipped. Default value is 0.5 - """ - - def __init__(self, p=0.5, with_coord=False): - self.p = p - self.with_coord = with_coord - - def __call__(self, results): - """ - Args: - img (PIL Image): Image to be flipped. - Returns: - PIL Image: Randomly flipped image. - """ - for key in results.get('img_fields', ['img']): - img = results[key] - if random.random() < self.p: - if self.with_coord: - coord = results['coord'] - coord_new = coord.clone() - coord_new[0] = coord[2] - coord_new[2] = coord[0] - results['coord'] = coord_new - img = F.hflip(img) - results[key] = img - return results - - def __repr__(self): - return self.__class__.__name__ + '(p={}, with_coord={})'.format(self.p, self.with_coord) - - -@PIPELINES.register_module() -class ColorJitter(object): - """Randomly change the brightness, contrast and saturation of an image. - - Args: - brightness (float or tuple of float (min, max)): How much to jitter brightness. - brightness_factor is chosen uniformly from [max(0, 1 - brightness), 1 + brightness] - or the given [min, max]. Should be non negative numbers. - contrast (float or tuple of float (min, max)): How much to jitter contrast. - contrast_factor is chosen uniformly from [max(0, 1 - contrast), 1 + contrast] - or the given [min, max]. Should be non negative numbers. - saturation (float or tuple of float (min, max)): How much to jitter saturation. - saturation_factor is chosen uniformly from [max(0, 1 - saturation), 1 + saturation] - or the given [min, max]. Should be non negative numbers. - hue (float or tuple of float (min, max)): How much to jitter hue. - hue_factor is chosen uniformly from [-hue, hue] or the given [min, max]. - Should have 0<= hue <= 0.5 or -0.5 <= min <= max <= 0.5. - """ - - def __init__(self, brightness=0, contrast=0, saturation=0, hue=0): - super().__init__() - self.brightness = self._check_input(brightness, 'brightness') - self.contrast = self._check_input(contrast, 'contrast') - self.saturation = self._check_input(saturation, 'saturation') - self.hue = self._check_input(hue, 'hue', center=0, bound=(-0.5, 0.5), - clip_first_on_zero=False) - - @torch.jit.unused - def _check_input(self, value, name, center=1, bound=(0, float('inf')), clip_first_on_zero=True): - if isinstance(value, numbers.Number): - if value < 0: - raise ValueError("If {} is a single number, it must be non negative.".format(name)) - value = [center - float(value), center + float(value)] - if clip_first_on_zero: - value[0] = max(value[0], 0.0) - elif isinstance(value, (tuple, list)) and len(value) == 2: - if not bound[0] <= value[0] <= value[1] <= bound[1]: - raise ValueError("{} values should be between {}".format(name, bound)) - else: - raise TypeError("{} should be a single number or a list/tuple with lenght 2.".format(name)) - - # if value is 0 or (1., 1.) for brightness/contrast/saturation - # or (0., 0.) for hue, do nothing - if value[0] == value[1] == center: - value = None - return value - - @staticmethod - @torch.jit.unused - def get_params(brightness, contrast, saturation, hue): - """Get a randomized transform to be applied on image. - - Arguments are same as that of __init__. - - Returns: - Transform which randomly adjusts brightness, contrast and - saturation in a random order. - """ - transforms = [] - - if brightness is not None: - brightness_factor = random.uniform(brightness[0], brightness[1]) - transforms.append(Lambda(lambda img: F.adjust_brightness(img, brightness_factor))) - - if contrast is not None: - contrast_factor = random.uniform(contrast[0], contrast[1]) - transforms.append(Lambda(lambda img: F.adjust_contrast(img, contrast_factor))) - - if saturation is not None: - saturation_factor = random.uniform(saturation[0], saturation[1]) - transforms.append(Lambda(lambda img: F.adjust_saturation(img, saturation_factor))) - - if hue is not None: - hue_factor = random.uniform(hue[0], hue[1]) - transforms.append(Lambda(lambda img: F.adjust_hue(img, hue_factor))) - - random.shuffle(transforms) - transform = Compose(transforms) - - return transform - - def __call__(self, results): - """ - Args: - img (PIL Image or Tensor): Input image. - - Returns: - PIL Image or Tensor: Color jittered image. - """ - for key in results.get('img_fields', ['img']): - img = results[key] - fn_idx = torch.randperm(4) - for fn_id in fn_idx: - if fn_id == 0 and self.brightness is not None: - brightness = self.brightness - brightness_factor = torch.tensor(1.0).uniform_(brightness[0], brightness[1]).item() - img = F.adjust_brightness(img, brightness_factor) - - if fn_id == 1 and self.contrast is not None: - contrast = self.contrast - contrast_factor = torch.tensor(1.0).uniform_(contrast[0], contrast[1]).item() - img = F.adjust_contrast(img, contrast_factor) - - if fn_id == 2 and self.saturation is not None: - saturation = self.saturation - saturation_factor = torch.tensor(1.0).uniform_(saturation[0], saturation[1]).item() - img = F.adjust_saturation(img, saturation_factor) - - if fn_id == 3 and self.hue is not None: - hue = self.hue - hue_factor = torch.tensor(1.0).uniform_(hue[0], hue[1]).item() - img = F.adjust_hue(img, hue_factor) - results[key] = img - return results - - def __repr__(self): - format_string = self.__class__.__name__ + '(' - format_string += 'brightness={0}'.format(self.brightness) - format_string += ', contrast={0}'.format(self.contrast) - format_string += ', saturation={0}'.format(self.saturation) - format_string += ', hue={0})'.format(self.hue) - return format_string - - -@PIPELINES.register_module() -class RandomGrayscale(object): - """Randomly convert image to grayscale with a probability of p (default 0.1). - The image can be a PIL Image or a Tensor, in which case it is expected - to have [..., 3, H, W] shape, where ... means an arbitrary number of leading - dimensions - - Args: - p (float): probability that image should be converted to grayscale. - - Returns: - PIL Image or Tensor: Grayscale version of the input image with probability p and unchanged - with probability (1-p). - - If input image is 1 channel: grayscale version is 1 channel - - If input image is 3 channel: grayscale version is 3 channel with r == g == b - - """ - - def __init__(self, p=0.1): - super().__init__() - self.p = p - - def __call__(self, results): - """ - Args: - img (PIL Image or Tensor): Image to be converted to grayscale. - - Returns: - PIL Image or Tensor: Randomly grayscaled image. - """ - for key in results.get('img_fields', ['img']): - img = results[key] - num_output_channels = F._get_image_num_channels(img) - if torch.rand(1) < self.p: - img = F.rgb_to_grayscale(img, num_output_channels=num_output_channels) - results[key] = img - return results - - def __repr__(self): - return self.__class__.__name__ + '(p={0})'.format(self.p) - - -@PIPELINES.register_module() -class GaussianBlur(object): - """Gaussian blur augmentation in SimCLR https://arxiv.org/abs/2002.05709.""" - - def __init__(self, sigma_min, sigma_max): - self.sigma_min = sigma_min - self.sigma_max = sigma_max - - def __call__(self, results): - for key in results.get('img_fields', ['img']): - img = results[key] - sigma = np.random.uniform(self.sigma_min, self.sigma_max) - img = img.filter(ImageFilter.GaussianBlur(radius=sigma)) - results[key] = img - return results - - def __repr__(self): - repr_str = self.__class__.__name__ - return repr_str - - -@PIPELINES.register_module() -class Solarization(object): - """Solarization augmentation in BYOL https://arxiv.org/abs/2006.07733.""" - - def __init__(self, threshold=128): - self.threshold = threshold - - def __call__(self, results): - for key in results.get('img_fields', ['img']): - img = results[key] - img = np.array(img) - img = np.where(img < self.threshold, img, 255 - img) - results[key] = Image.fromarray(img.astype(np.uint8)) - return results - - def __repr__(self): - repr_str = self.__class__.__name__ - return repr_str - - -@PIPELINES.register_module() -class RandomAppliedTrans(object): - """Randomly applied transformations. - - Args: - transforms (list[dict]): List of transformations in dictionaries. - p (float): Probability. - """ - - def __init__(self, transforms, p=0.5): - t = [build_pipeline(t) for t in transforms] - self.trans = T.RandomApply(t, p=p) - - def __call__(self, results): - return self.trans(results) - - def __repr__(self): - repr_str = self.__class__.__name__ - return repr_str - - -@PIPELINES.register_module() -class Normalize(object): - """Normalize the image. - - Args: - mean (sequence): Mean values of 3 channels. - std (sequence): Std values of 3 channels. - to_rgb (bool): Whether to convert the image from BGR to RGB, - default is true. - """ - - def __init__(self, mean, std, to_rgb=True): - self.mean = np.array(mean, dtype=np.float32) - self.std = np.array(std, dtype=np.float32) - self.to_rgb = to_rgb - - def __call__(self, results): - for key in results.get('img_fields', ['img']): - results[key] = mmcv.imnormalize(results[key], self.mean, self.std, - self.to_rgb) - results['img_norm_cfg'] = dict( - mean=self.mean, std=self.std, to_rgb=self.to_rgb) - return results - - def __repr__(self): - repr_str = self.__class__.__name__ - repr_str += f'(mean={list(self.mean)}, ' - repr_str += f'std={list(self.std)}, ' - repr_str += f'to_rgb={self.to_rgb})' - return repr_str - - -@PIPELINES.register_module() -class ToNumpy(object): - def __init__(self): - pass - - def __call__(self, results): - for key in results.get('img_fields', ['img']): - results[key] = np.asarray(results[key], dtype=np.float32) - return results - - -@PIPELINES.register_module() -class ImageToTensor(object): - - def __init__(self, keys): - self.keys = keys - - def __call__(self, results): - for key in self.keys: - img = results[key] - if len(img.shape) < 3: - img = np.expand_dims(img, -1) - results[key] = to_tensor(img.transpose(2, 0, 1)) - return results - - def __repr__(self): - return self.__class__.__name__ + f'(keys={self.keys})' - - -@PIPELINES.register_module() -class Collect(object): - """ - Collect data from the loader relevant to the specific task. - - This is usually the last stage of the data loader pipeline. Typically keys - is set to some subset of "img" and "gt_label". - - Args: - keys (Sequence[str]): Keys of results to be collected in ``data``. - meta_keys (Sequence[str], optional): Meta keys to be converted to - ``mmcv.DataContainer`` and collected in ``data[img_metas]``. - Default: ``('filename', 'ori_shape', 'img_shape', 'flip', - 'flip_direction', 'img_norm_cfg')`` - - Returns: - dict: The result dict contains the following keys - - keys in``self.keys`` - - ``img_metas`` if avaliable - """ - - def __init__(self, - keys, - meta_keys=('filename', 'ori_filename', 'ori_shape', - 'img_shape', 'flip', 'flip_direction', - 'img_norm_cfg')): - self.keys = keys - self.meta_keys = meta_keys - - def __call__(self, results): - data = {} - img_meta = {} - for key in self.meta_keys: - if key in results: - img_meta[key] = results[key] - data['img_metas'] = DC(img_meta, cpu_only=True) - for key in self.keys: - data[key] = results[key] - return data - - def __repr__(self): - return self.__class__.__name__ + \ - f'(keys={self.keys}, meta_keys={self.meta_keys})' diff --git a/mpa/modules/datasets/selfsl_dataset.py b/mpa/modules/datasets/selfsl_dataset.py deleted file mode 100644 index ac5b17ea..00000000 --- a/mpa/modules/datasets/selfsl_dataset.py +++ /dev/null @@ -1,73 +0,0 @@ -import os.path as osp -import numpy as np -import PIL.Image as Image -import torch -from torch.utils.data import Dataset -from torchvision.transforms import Compose - -from mpa.selfsl.builder import DATASETS -from mpa.selfsl.builder import build_datasource, build_pipeline - - -@DATASETS.register_module() -class SelfSLDataset(Dataset): - """Wrapper of existing dataset for SelfSL - """ - def __init__(self, down_task, datasource, pipeline): - self.datasource = build_datasource(datasource, down_task) - - pipeline1 = [build_pipeline(p) for p in pipeline['view0']] - self.pipeline1 = Compose(pipeline1) - pipeline2 = [build_pipeline(p) for p in pipeline['view1']] - self.pipeline2 = Compose(pipeline2) - - def __len__(self): - return len(self.datasource) - - def _getitem_from_source(self, idx): - data = self.datasource[idx] - if isinstance(data, dict): # dict(img:ndarray) - if 'img' not in data.keys(): - if data['img_prefix'] is not None: - filename = osp.join(data['img_prefix'], - data['img_info']['filename']) - else: - filename = data['img_info']['filename'] - data['filename'] = filename - data['img'] = Image.open(filename) - img = data['img'] - elif isinstance(data, tuple): # (img:[ndarray|Image], gt) - img, _ = data - else: - img = data - - if isinstance(img, np.ndarray): - img = Image.fromarray(img) - elif not isinstance(img, Image.Image): - raise TypeError('image should be PIL.Image or numpy.ndarray \ - - {}'.format(type(img))) - - if img.mode != 'RGB': - img = img.convert('RGB') - - return img - - def __getitem__(self, idx): - if torch.is_tensor(idx): - idx = idx.tolist() - - img = self._getitem_from_source(idx) - - results1 = self.pipeline1(dict(img=img)) - results2 = self.pipeline2(dict(img=img)) - - results = dict() - for k, v in results1.items(): - results[k+'1'] = v - for k, v in results2.items(): - results[k+'2'] = v - - return results - - def evaluate(self, **kwargs): - pass diff --git a/mpa/modules/models/__init__.py b/mpa/modules/models/__init__.py index f5877fd9..2022a508 100644 --- a/mpa/modules/models/__init__.py +++ b/mpa/modules/models/__init__.py @@ -6,7 +6,5 @@ from . import classifiers from . import detectors -from . import selfsl_model - from . import segmentors -from .scalar_schedulers import * \ No newline at end of file +from .scalar_schedulers import * diff --git a/mpa/modules/models/heads/__init__.py b/mpa/modules/models/heads/__init__.py index a36f7952..f7933e59 100644 --- a/mpa/modules/models/heads/__init__.py +++ b/mpa/modules/models/heads/__init__.py @@ -4,7 +4,6 @@ from . import task_incremental_classifier_head from . import non_linear_cls_head from . import cls_incremental_head -from . import selfsl_head from . import custom_vfnet_head from . import custom_atss_head from . import custom_retina_head diff --git a/mpa/modules/models/heads/selfsl_head.py b/mpa/modules/models/heads/selfsl_head.py deleted file mode 100644 index 8a42762f..00000000 --- a/mpa/modules/models/heads/selfsl_head.py +++ /dev/null @@ -1,96 +0,0 @@ -import torch -import torch.nn as nn - -from mpa.selfsl.builder import HEADS -from mpa.selfsl.builder import build_neck - - -def MSELoss(x, y): - loss = 2 * x.size(0) - 2 * (x * y).sum() - return loss / x.size(0) - - -def PPCLoss(q, k, coord_q, coord_k, pos_ratio): - """ Pixel to Propagation Consistency Loss - q, k: N * C * H * W - coord_q, coord_k: N * 4 (x_upper_left, y_upper_left, x_lower_right, y_lower_right) - """ - N, C, H, W = q.shape - # [bs, feat_dim, 49] - q = q.view(N, C, -1) - k = k.view(N, C, -1) - - # generate center_coord, width, height - # [1, 7, 7] - x_array = torch.arange(0., float(W), dtype=coord_q.dtype, device=coord_q.device).view(1, 1, -1).repeat(1, H, 1) - y_array = torch.arange(0., float(H), dtype=coord_q.dtype, device=coord_q.device).view(1, -1, 1).repeat(1, 1, W) - # [bs, 1, 1] - q_bin_width = ((coord_q[:, 2] - coord_q[:, 0]) / W).view(-1, 1, 1) - q_bin_height = ((coord_q[:, 3] - coord_q[:, 1]) / H).view(-1, 1, 1) - k_bin_width = ((coord_k[:, 2] - coord_k[:, 0]) / W).view(-1, 1, 1) - k_bin_height = ((coord_k[:, 3] - coord_k[:, 1]) / H).view(-1, 1, 1) - # [bs, 1, 1] - q_start_x = coord_q[:, 0].view(-1, 1, 1) - q_start_y = coord_q[:, 1].view(-1, 1, 1) - k_start_x = coord_k[:, 0].view(-1, 1, 1) - k_start_y = coord_k[:, 1].view(-1, 1, 1) - - # [bs, 1, 1] - q_bin_diag = torch.sqrt(q_bin_width ** 2 + q_bin_height ** 2) - k_bin_diag = torch.sqrt(k_bin_width ** 2 + k_bin_height ** 2) - max_bin_diag = torch.max(q_bin_diag, k_bin_diag) - - # [bs, 7, 7] - center_q_x = (x_array + 0.5) * q_bin_width + q_start_x - center_q_y = (y_array + 0.5) * q_bin_height + q_start_y - center_k_x = (x_array + 0.5) * k_bin_width + k_start_x - center_k_y = (y_array + 0.5) * k_bin_height + k_start_y - - # [bs, 49, 49] - dist_center = torch.sqrt((center_q_x.view(-1, H * W, 1) - center_k_x.view(-1, 1, H * W)) ** 2 - + (center_q_y.view(-1, H * W, 1) - center_k_y.view(-1, 1, H * W)) ** 2) / max_bin_diag - pos_mask = (dist_center < pos_ratio).float().detach() - - # [bs, 49, 49] - logit = torch.bmm(q.transpose(1, 2), k) - - loss = (logit * pos_mask).sum(-1).sum(-1) / (pos_mask.sum(-1).sum(-1) + 1e-6) - - return -2 * loss.mean() - - -@HEADS.register_module() -class LatentPredictHead(nn.Module): - """Head for contrastive learning. - """ - - def __init__(self, predictor, loss, **kwargs): - super(LatentPredictHead, self).__init__() - self.predictor = build_neck(predictor) - if loss == 'MSE': - self.regression_loss = MSELoss - elif loss == 'PPC': - pos_ratio = kwargs.get('pos_ratio', 0.5) - self.regression_loss = lambda *args: PPCLoss(*args, pos_ratio=pos_ratio) - else: - raise ValueError(f'{loss} loss is not supported') - - def init_weights(self, init_linear='normal'): - self.predictor.init_weights(init_linear=init_linear) - - def forward(self, proj, proj_ng, *args): - """Forward head. - - Args: - input (Tensor): NxC input features. - target (Tensor): NxC target features. - - Returns: - dict[str, Tensor]: A dictionary of loss components. - """ - pred = self.predictor(proj) - pred_norm = nn.functional.normalize(pred, dim=1) - proj_ng_norm = nn.functional.normalize(proj_ng, dim=1) - - loss = self.regression_loss(pred_norm, proj_ng_norm, *args) - return dict(loss=loss) diff --git a/mpa/modules/models/necks/__init__.py b/mpa/modules/models/necks/__init__.py index 7038fe64..9c0fa90a 100644 --- a/mpa/modules/models/necks/__init__.py +++ b/mpa/modules/models/necks/__init__.py @@ -1,2 +1 @@ # flake8: noqa -from . import selfsl_neck \ No newline at end of file diff --git a/mpa/modules/models/necks/selfsl_neck.py b/mpa/modules/models/necks/selfsl_neck.py deleted file mode 100644 index 595b17e0..00000000 --- a/mpa/modules/models/necks/selfsl_neck.py +++ /dev/null @@ -1,118 +0,0 @@ -import torch -import torch.nn as nn -import torch.nn.functional as F -from mmcv.cnn import kaiming_init, normal_init -from mmcv.cnn import build_norm_layer - -from mpa.selfsl.builder import NECKS - - -def _init_weights(module, init_linear='normal', std=0.01, bias=0.): - if init_linear not in ['normal', 'kaiming']: - raise ValueError("Undefined init_linear: {}".format(init_linear)) - for m in module.modules(): - if isinstance(m, nn.Linear): - if init_linear == 'normal': - normal_init(m, std=std, bias=bias) - else: - kaiming_init(m, mode='fan_in', nonlinearity='relu') - elif isinstance(m, (nn.BatchNorm1d, nn.BatchNorm2d, - nn.GroupNorm, nn.SyncBatchNorm)): - if m.weight is not None: - nn.init.constant_(m.weight, 1) - if m.bias is not None: - nn.init.constant_(m.bias, 0) - - -@NECKS.register_module() -class MLP(nn.Module): - """The MLP neck: fc/conv-bn-relu-fc/conv. - """ - - def __init__(self, - in_channels, - hid_channels, - out_channels, - norm_cfg=dict(type='BN1d'), - use_conv=False, - with_avg_pool=True): - super(MLP, self).__init__() - - self.with_avg_pool = with_avg_pool - if with_avg_pool: - self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) - - self.use_conv = use_conv - if use_conv: - self.mlp = nn.Sequential( - nn.Conv2d(in_channels, hid_channels, 1), - build_norm_layer(norm_cfg, hid_channels)[1], - nn.ReLU(inplace=True), - nn.Conv2d(hid_channels, out_channels, 1) - ) - else: - self.mlp = nn.Sequential( - nn.Linear(in_channels, hid_channels), - build_norm_layer(norm_cfg, hid_channels)[1], - nn.ReLU(inplace=True), - nn.Linear(hid_channels, out_channels) - ) - - def init_weights(self, init_linear='normal'): - _init_weights(self, init_linear) - - def forward(self, x): - if isinstance(x, (tuple, list)): - # using last output - x = x[-1] - if not isinstance(x, torch.Tensor): - raise TypeError('neck inputs should be tuple or torch.tensor') - if self.with_avg_pool: - x = self.avgpool(x) - if self.use_conv: - return self.mlp(x) - else: - return self.mlp(x.view(x.size(0), -1)) - - -@NECKS.register_module() -class PPM(nn.Module): - """ - Head for Pixpro - """ - def __init__(self, sharpness, **kwargs): - super(PPM, self).__init__() - self.transform = nn.Conv2d(256, 256, 1) - self.sharpness = sharpness - - def init_weights(self, **kwargs): - pass - - def forward(self, feat): - if isinstance(feat, (tuple, list)): - feat = feat[-1] - if not isinstance(feat, torch.Tensor): - raise TypeError('neck inputs should be tuple or torch.tensor') - - N, C, H, W = feat.shape - - # Value transformation - feat_value = self.transform(feat) - feat_value = F.normalize(feat_value, dim=1) - feat_value = feat_value.view(N, C, -1) - - # Similarity calculation - feat = F.normalize(feat, dim=1) - - # [N, C, H * W] - feat = feat.view(N, C, -1) - - # [N, H * W, H * W] - attention = torch.bmm(feat.transpose(1, 2), feat) - attention = torch.clamp(attention, min=0.) - attention = attention ** self.sharpness - - # [N, C, H * W] - feat = torch.bmm(feat_value, attention.transpose(1, 2)) - - return feat.view(N, C, H, W) diff --git a/mpa/modules/models/selfsl_model.py b/mpa/modules/models/selfsl_model.py deleted file mode 100755 index 6940f4ac..00000000 --- a/mpa/modules/models/selfsl_model.py +++ /dev/null @@ -1,187 +0,0 @@ -from collections import OrderedDict -import torch -import torch.nn as nn -import torch.distributed as dist - -from mpa.selfsl.builder import TRAINERS -from mpa.selfsl.builder import build_backbone, build_head, build_neck - - -@TRAINERS.register_module() -class SelfSL(nn.Module): - """SelfSL - BYOL/PixPro - - Args: - down_task (str): Type of downstream task. - backbone (dict): Config dict for module of backbone ConvNet. - neck (dict): Config dict for module of deep features to compact feature vectors. - Default: None. - head (dict): Config dict for module of loss functions. Default: None. - pretrained (str, optional): Path to pre-trained weights. Default: None. - base_momentum (float): The base momentum coefficient for the target network. - Default: 0.996. - """ - - def __init__(self, - down_task, - backbone, - neck=None, - head=None, - pretrained=None, - base_momentum=0.996, - **kwargs): - super(SelfSL, self).__init__() - - # build backbone - self.online_backbone = build_backbone(backbone, down_task) - self.target_backbone = build_backbone(backbone, down_task) - - # build projector - self.online_projector = build_neck(neck) - self.target_projector = build_neck(neck) - - # build head with predictor - self.head = build_head(head) - - self.init_weights(pretrained=pretrained) - - self.base_momentum = base_momentum - self.momentum = base_momentum - - def init_weights(self, pretrained=None): - """Initialize the weights of model. - - Args: - pretrained (str, optional): Path to pre-trained weights. - Default: None. - """ - if pretrained is not None: - print('load model from: {}'.format(pretrained)) - - # init backbone - self.online_backbone.init_weights(pretrained=pretrained) - for param_ol, param_tgt in zip(self.online_backbone.parameters(), - self.target_backbone.parameters()): - param_tgt.data.copy_(param_ol.data) - param_tgt.requires_grad = False - param_ol.requires_grad = True - - # init projector - self.online_projector.init_weights(init_linear='kaiming') - for param_ol, param_tgt in zip(self.online_projector.parameters(), - self.target_projector.parameters()): - param_tgt.data.copy_(param_ol.data) - param_tgt.requires_grad = False - param_ol.requires_grad = True - - # init the predictor - self.head.init_weights() - - @torch.no_grad() - def _momentum_update(self): - """Momentum update of the target network.""" - for param_ol, param_tgt in zip(self.online_backbone.parameters(), - self.target_backbone.parameters()): - param_tgt.data = param_tgt.data * self.momentum + \ - param_ol.data * (1. - self.momentum) - - for param_ol, param_tgt in zip(self.online_projector.parameters(), - self.target_projector.parameters()): - param_tgt.data = param_tgt.data * self.momentum + \ - param_ol.data * (1. - self.momentum) - - @torch.no_grad() - def momentum_update(self): - self._momentum_update() - - def forward(self, img1, img2, **kwargs): - """Forward computation during training. - - Args: - img (Tensor): Input of two concatenated images of shape (N, 2, C, H, W). - Typically these should be mean centered and std scaled. - - Returns: - dict[str, Tensor]: A dictionary of loss components. - """ - proj_1 = self.online_projector(self.online_backbone(img1)) - proj_2 = self.online_projector(self.online_backbone(img2)) - with torch.no_grad(): - proj_1_tgt = self.target_projector(self.target_backbone(img1)).clone().detach() - proj_2_tgt = self.target_projector(self.target_backbone(img2)).clone().detach() - - coord_1_2 = [] - coord_2_1 = [] - if 'coord1' in kwargs and 'coord2' in kwargs: - coord1 = kwargs['coord1'] - coord2 = kwargs['coord2'] - coord_1_2 = [coord1, coord2] - coord_2_1 = [coord2, coord1] - - loss = self.head(proj_1, proj_2_tgt, *coord_1_2)['loss'] \ - + self.head(proj_2, proj_1_tgt, *coord_2_1)['loss'] - return dict(loss=loss) - - def train_step(self, data, optimizer): - """The iteration step during training. - - This method defines an iteration step during training, except for the - back propagation and optimizer updating, which are done in an optimizer - hook. Note that in some complicated cases or models, the whole process - including back propagation and optimizer updating are also defined in - this method, such as GAN. - - Args: - data (dict): The output of dataloader. - optimizer (:obj:`torch.optim.Optimizer` | dict): The optimizer of - runner is passed to ``train_step()``. This argument is unused - and reserved. - - Returns: - dict: It should contain at least 3 keys: ``loss``, ``log_vars``, - ``num_samples``. - ``loss`` is a tensor for back propagation, which can be a - weighted sum of multiple losses. - ``log_vars`` contains all the variables to be sent to the - logger. - ``num_samples`` indicates the batch size (when the model is - DDP, it means the batch size on each GPU), which is used for - averaging the logs. - """ - losses = self(**data) - loss, log_vars = self._parse_losses(losses) - - outputs = dict( - loss=loss, log_vars=log_vars, num_samples=len(data['img1'].data)) - - return outputs - - def val_step(self, *args): - pass - - def _parse_losses(self, losses): - log_vars = OrderedDict() - for loss_name, loss_value in losses.items(): - if isinstance(loss_value, torch.Tensor): - log_vars[loss_name] = loss_value.mean() - elif isinstance(loss_value, list): - log_vars[loss_name] = sum(_loss.mean() for _loss in loss_value) - elif isinstance(loss_value, dict): - for name, value in loss_value.items(): - log_vars[name] = value - else: - raise TypeError( - f'{loss_name} is not a tensor or list of tensors') - - loss = sum(_value for _key, _value in log_vars.items() - if 'loss' in _key) - - log_vars['loss'] = loss - for loss_name, loss_value in log_vars.items(): - # reduce loss when distributed training - if dist.is_available() and dist.is_initialized(): - loss_value = loss_value.data.clone() - dist.all_reduce(loss_value.div_(dist.get_world_size())) - log_vars[loss_name] = loss_value.item() - - return loss, log_vars diff --git a/mpa/selfsl/__init__.py b/mpa/selfsl/__init__.py deleted file mode 100755 index 50c97521..00000000 --- a/mpa/selfsl/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# flake8: noqa -from . import trainer \ No newline at end of file diff --git a/mpa/selfsl/builder.py b/mpa/selfsl/builder.py deleted file mode 100644 index e559e10c..00000000 --- a/mpa/selfsl/builder.py +++ /dev/null @@ -1,142 +0,0 @@ -import numpy as np -import random -from functools import partial -from torch import nn -from torch.utils.data import DataLoader -from torch.utils.data.dataset import ConcatDataset -from mmcv.parallel import collate -from mmcv.runner import get_dist_info -from mmcv.utils import build_from_cfg -from mmcv.utils import Registry - -from mpa.modules.datasets.samplers.distributed_sampler import DistributedSampler - -TRAINERS = Registry('trainers') -HEADS = Registry('heads') -NECKS = Registry('necks') -DATASETS = Registry('datasets') -PIPELINES = Registry('pipeline') - - -def build(cfg, registry, default_args=None): - if isinstance(cfg, list): - modules = [ - build_from_cfg(cfg_, registry, default_args) for cfg_ in cfg - ] - return nn.Sequential(*modules) - else: - return build_from_cfg(cfg, registry, default_args) - - -def build_backbone(cfg, down_task): - if down_task == 'classification': - from mmcls.models import BACKBONES - elif down_task == 'detection': - from mmdet.models import BACKBONES - return build(cfg, BACKBONES) - - -def build_head(cfg): - return build(cfg, HEADS) - - -def build_neck(cfg): - return build(cfg, NECKS) - - -def build_trainer(cfg): - return build(cfg, TRAINERS) - - -def build_pipeline(cfg): - return build_from_cfg(cfg, PIPELINES) - - -def build_datasource(cfg, down_task, default_args=None): - if down_task == 'classification': - from mmcls.datasets import DATASETS as DATASOURCES - elif down_task == 'detection': - from mmdet.datasets import DATASETS as DATASOURCES - - if isinstance(cfg, (list, tuple)): - dataset = ConcatDataset([build_datasource(c, down_task, default_args) for c in cfg]) - else: - dataset = build_from_cfg(cfg, DATASOURCES, default_args) - - return dataset - - -def build_dataset(cfg, default_args=None): - if isinstance(cfg, (list, tuple)): - dataset = ConcatDataset([build_dataset(c, default_args) for c in cfg]) - else: - dataset = build_from_cfg(cfg, DATASETS, default_args) - - return dataset - - -def build_dataloader(dataset, - samples_per_gpu, - workers_per_gpu, - num_gpus=1, - dist=True, - shuffle=True, - seed=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. - kwargs: any keyword argument to be used to initialize DataLoader - - Returns: - DataLoader: A PyTorch dataloader. - """ - rank, world_size = get_dist_info() - if dist: - sampler = DistributedSampler(dataset, world_size, rank, shuffle) - shuffle = False - batch_size = samples_per_gpu - num_workers = workers_per_gpu - else: - sampler = None - 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 - - 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=False, - 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) diff --git a/mpa/selfsl/trainer.py b/mpa/selfsl/trainer.py deleted file mode 100644 index 5107321d..00000000 --- a/mpa/selfsl/trainer.py +++ /dev/null @@ -1,391 +0,0 @@ -import numbers -import time -import torch -import os -import os.path as osp -import warnings - -import torch.multiprocessing as mp -import torch.distributed as dist - -from mmcv import __version__ -from mmcv import collect_env -from mmcv import get_git_hash -from mmcv.parallel import MMDataParallel, MMDistributedDataParallel -from mmcv.runner import build_optimizer, build_runner -from mmcv.runner import load_checkpoint -from mmcv.runner import DistSamplerSeedHook - -from mmcls import __version__ as mmcls_version -from mmdet import __version__ as mmdet_version -from mmdet.parallel import MMDataCPU - -from mpa.registry import STAGES -from mpa.stage import Stage -from mpa.stage import _set_random_seed -from mpa.utils.convert_keys import convert_keys -from mpa.utils.logger import get_logger - -from .builder import build_backbone, build_trainer, build_dataset, build_dataloader - -logger = get_logger() - - -@STAGES.register_module() -class SelfSLTrainer(Stage): - """ - Run a self supervised training with given backbone and dataset without annotation. - The default selfsl method is 'BYOL' and base configuration file is provided by deafult. - - Args: - cfg (str, optional): Path to config file of SelfSL. - downstream_cfg (str, optional): Path to config file of downstream task. - downstream_task (str, optional): Type of downstream task, e.g., classification. - """ - def __init__(self, **kwargs): - super().__init__(**kwargs) - self.down_task = self.cfg.pop('down_task') - self.distributed = False - - def _configure_data(self, data_cfg): - """ - Configure dataset for SelfSL. - The given dataset will be a dataource in dataset wrapper for SelfSL. - - Args: - data_cfg (Config): An dataset config as datasource. - """ - cfg = self.cfg - cfg.data.train.down_task = self.down_task - - # get and put pipeline options from data_cfg to cfg. it will be used by Config.configure_data() - # to update pipelines - poptions = data_cfg['data'].get('pipeline_options') - if poptions is not None: - cfg.data['pipeline_options'] = poptions - - # TODO: enable to use data.train? - datasource_cfg = data_cfg['data']['unlabeled'] - - # Datasource config - if not isinstance(datasource_cfg, (list, tuple)): - datasource_cfg = [datasource_cfg] - - # remove pipelines in datasource - for _cfg in datasource_cfg: - _cfg['pipeline'] = [] - if self.down_task == 'detection': - _cfg.filter_empty_gt = False - - cfg.data.train.datasource = datasource_cfg - Stage.configure_data(cfg, True) - - return cfg - - def _configure_model(self, backbone_cfg, model_ckpt): - """ - Configure model for SelfSL with the given backbone cfg. - If the pretrained is NoneType, the ImageNet pretrained model provided by torch or mmcls - will be loaded by default. - - Args: - backbone_cfg (Config): A backbone config. - model_ckpt (str): Path to pretrained model of backbone for warm starts. - Set 'random' to train from scratch. - """ - cfg = self.cfg - cfg.model.down_task = self.down_task - - # Load imagenet pretrained model by default for warm-start - # resnet18/34/50/101/152, mobilenet_v2, ote_mobilenetv3_small/large/large_075, ote_efficientnet_b0 - # TODO: changed default init weight t2, ote_mobilenetv3_small/large/large_075, ote_efficientne - backbone_type = backbone_cfg.type.lower() - if 'ote' in backbone_type: - cfg.model.pretrained = True - elif 'resnet' in backbone_type: - depth = backbone_cfg.depth - if depth in [18, 34, 50, 101, 152]: - cfg.model.pretrained = 'torchvision://resnet{}'.format(depth) - elif 'mobilenet' in backbone_type: - if 'v2' in backbone_type: - cfg.model.pretrained = 'mmcls://mobilenet_v2' - - if not cfg.model.pretrained: - logger.info(f'no pretrained model for {backbone_type}') - - if model_ckpt == 'random': - logger.info('train from scratch') - cfg.model.pretrained = None - elif model_ckpt: - if 'ote' in backbone_type: - new_path = osp.join(cfg.work_dir, osp.basename(model_ckpt)[:-3] + 'converted.pth')\ - if cfg.work_dir else None - cfg.load_from = convert_keys(backbone_cfg.type, model_ckpt, new_path) - else: - cfg.load_from = model_ckpt - cfg.model.pretrained = None - - # get the number of output channles of backbone - backbone = build_backbone(backbone_cfg, self.down_task) - output = backbone(torch.rand([1, 3, 224, 224])) - if isinstance(output, (tuple, list)): - output = output[-1] - out_channels = output.shape[1] - - # set in_channel of neck as out_channel of backbone - if 'neck' in cfg.model: - cfg.model.neck.in_channels = out_channels - - # unfreeze backbone and bn layers - backbone_cfg['frozen_stages'] = -1 - if backbone_cfg.get('norm_eval'): - backbone_cfg.norm_eval = False - - # change to syncBN in multi-gpu mode - if self.distributed: - # TODO: need an exeption handling for models w/o BN - if 'resnet' in backbone_cfg.type: - backbone_cfg.norm_cfg = dict(type='SyncBN') - cfg.model.neck.norm_cfg = dict(type='SyncBN') - cfg.model.head.predictor.norm_cfg = dict(type='SyncBN') - - cfg.model.backbone = backbone_cfg - return cfg - - def configure(self, backbone_cfg, data_cfg, model_ckpt=None, resume=None): - """ - Configure overall SelfSL pipeline with given backbone and dataset. - - Args: - backbone_cfg (Config): A backbone config. - data_cfg (Config): An dataset config. - model_ckpt (str, optional): Path to pretrained model of backbone. - Set 'random' to train from scratch. - resume (str, optional): Path to model to resume from. - """ - logger.info('called configure selfsl trainer') - cfg = self.cfg - - # Distributed - if cfg.gpu_ids is not None: - if isinstance(cfg.get('gpu_ids'), numbers.Number): - cfg.gpu_ids = [cfg.get('gpu_ids')] - if len(cfg.gpu_ids) > 1: - self.distributed = True - - # Environment - if not hasattr(cfg, 'deterministic'): - cfg.deterministic = True - _set_random_seed(cfg.seed, cfg.deterministic) - logger.info(f'set random seed:{cfg.seed}, deterministic:{cfg.deterministic}') - - # Model - self._configure_model(backbone_cfg, model_ckpt) - - if resume: - cfg.resume_from = resume - - # Data - self._configure_data(data_cfg) - - # add 'num_gpus' and 'batch_size' to the cfg - if self.distributed: - cfg.num_gpus = len(cfg.gpu_ids) - cfg.batch_size = cfg.num_gpus * cfg.data.samples_per_gpu - else: - if torch.cuda.is_available(): - cfg.num_gpus = 1 - else: - cfg.num_gpus = 0 - cfg.batch_size = cfg.data.samples_per_gpu - - # applying linear scaling rule - cfg.base_lr = cfg.optimizer.lr - cfg.optimizer.lr *= (cfg.batch_size / 256) - logger.info(f'batch size is {cfg.batch_size} ({cfg.data.samples_per_gpu} x {max(cfg.num_gpus, 1)} gpus)') - logger.info(f'scaling lr from {cfg.base_lr} to {cfg.optimizer.lr} following linear scaling rule') - - # TODO: only apply linear scaling rule if 'dist_params.linear_scale_lr' is set ??? - # if self.distributed: - # if cfg.dist_params.get('linear_scale_lr', False): - # new_lr = len(cfg.gpu_ids) * cfg.optimizer.lr - # logger.info(f'enabled linear scaling rule to the learning rate. \ - # changed LR from {cfg.optimizer.lr} to {new_lr}') - # cfg.optimizer.lr = new_lr - - return cfg - - def run(self, model_cfg, model_ckpt, data_cfg, **kwargs): - """Run training stage for selfsl - - - Configuration - - Environment setup - - Run training - - Save backbone model as 'backbone.pth' - """ - # Init logger - timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime()) - self._init_logger() - logger.info(f'called run with {kwargs}') - mode = kwargs.get('mode', 'train') - if mode not in self.mode: - return {} - - down_task = model_cfg.get('task', None) - if down_task and down_task != self.down_task: - logger.error(f'wrong {down_task} model for {self.down_task} task') - raise ValueError() - - # Configure - cfg = self.configure(model_cfg.model.backbone, data_cfg, model_ckpt) - - # Environment - 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) - - # Metadata - meta = dict() - meta['env_info'] = env_info - meta['config'] = cfg.pretty_text - meta['seed'] = cfg.seed - meta['exp_name'] = cfg.model.type + cfg.work_dir - if cfg.checkpoint_config is not None: - cfg.checkpoint_config.meta = dict( - mmcv_version=__version__ + '.' + get_git_hash()[:7], - mmdet_version=mmdet_version, - mmcls_version=mmcls_version) - - # Save config - cfg.dump(osp.join(cfg.work_dir, 'config.yaml')) - logger.info(f'Config:\n{cfg.pretty_text}') - - # Data - datasets = [build_dataset(cfg.data.train)] - - if self.distributed: - os.environ['MASTER_ADDR'] = cfg.dist_params.get('master_addr', 'localhost') - os.environ['MASTER_PORT'] = str(cfg.dist_params.get('master_port', '29500')) - mp.spawn(SelfSLTrainer.train_worker, nprocs=len(cfg.gpu_ids), - args=(datasets, cfg, self.distributed, timestamp, meta)) - else: - SelfSLTrainer.train_worker(None, datasets, cfg, - distributed=self.distributed, - timestamp=timestamp, - meta=meta) - - # Extract backbone weights - chk = torch.load(osp.join(cfg.work_dir, 'latest.pth')) - backbone_chk = dict(meta=chk['meta'], optimizer=chk['optimizer'], state_dict=dict()) - for k, v in chk['state_dict'].items(): - if k.startswith('online_backbone'): - if 'OTE' in cfg.model.backbone.type: - backbone_chk['state_dict'][k[16:]] = v - else: - backbone_chk['state_dict'][k[7:]] = v - - if len(backbone_chk['state_dict']) == 0: - raise Exception('Cannot find a backbone module in the checkpoint') - - backbone_ckpt_path = osp.join(cfg.work_dir, 'backbone.pth') - torch.save(backbone_chk, backbone_ckpt_path) - - return dict(pretrained=backbone_ckpt_path) - - @staticmethod - def train_worker(gpu, - dataset, - cfg, - distributed=False, - timestamp=None, - meta=None): - if distributed: - torch.cuda.set_device(gpu) - dist.init_process_group(backend=cfg.dist_params.get('backend', 'nccl'), - world_size=len(cfg.gpu_ids), rank=gpu) - logger.info(f'dist info world_size = {dist.get_world_size()}, rank = {dist.get_rank()}') - - # Model - model = build_trainer(cfg.model) - - # prepare data loaders - dataset = dataset if isinstance(dataset, (list, tuple)) else [dataset] - - data_loaders = [ - build_dataloader( - ds, - cfg.data.samples_per_gpu, - cfg.data.workers_per_gpu, - dist=distributed, - seed=cfg.seed) for ds in dataset - ] - - # put model on gpus - if torch.cuda.is_available(): - if distributed: - find_unused_parameters = cfg.get('find_unused_parameters', False) - model = MMDistributedDataParallel( - model.cuda(), - device_ids=[torch.cuda.current_device()], - broadcast_buffers=False, - find_unused_parameters=find_unused_parameters) - else: - model = MMDataParallel( - model.cuda(), device_ids=[torch.cuda.current_device()]) - else: - model = MMDataCPU(model) - - # build runner - optimizer = build_optimizer(model, cfg.optimizer) - - if cfg.get('runner') is None: - cfg.runner = { - 'type': 'EpochBasedRunner', - 'max_epochs': cfg.total_epochs - } - warnings.warn( - 'config is now expected to have a `runner` section, ' - 'please set `runner` in your config.', UserWarning) - - runner = build_runner( - cfg.runner, - default_args=dict( - model=model, - batch_processor=None, - optimizer=optimizer, - work_dir=cfg.work_dir, - logger=logger, - meta=meta)) - - # an ugly walkaround to make the .log and .log.json filenames the same - runner.timestamp = timestamp - - # register hooks - runner.register_training_hooks(cfg.lr_config, cfg.optimizer_config, - cfg.checkpoint_config, cfg.log_config, - cfg.get('momentum_config', None)) - - # When distributed training, it is only useful in conjunction with 'EpochBasedRunner`, - # while `IterBasedRunner` achieves the same purpose with `IterLoader`. - if distributed and cfg.runner.type == 'EpochBasedRunner': - runner.register_hook(DistSamplerSeedHook()) - - for hook in cfg.get('custom_hooks', ()): - runner.register_hook_from_cfg(hook) - - if cfg.get('resume_from', False): - runner.resume(cfg.resume_from) - elif cfg.get('load_from', False): - logger.info(f'load checkpoint from {cfg.load_from}') - try: - # workaround code to reduce a bunch of logs about missing keys during loading checkpoint - _backbone = torch.nn.Module() - _backbone.add_module( - 'backbone', - getattr(runner.model, 'module', runner.model).online_backbone) - load_checkpoint(_backbone, cfg.load_from, map_location='cpu', strict=False, logger=logger) - except IOError: - logger.warn('cannot open {}'.format(cfg.load_from)) - - runner.run(data_loaders, cfg.workflow) diff --git a/mpa/utils/hpo_stage.py b/mpa/utils/hpo_stage.py deleted file mode 100644 index abb27baa..00000000 --- a/mpa/utils/hpo_stage.py +++ /dev/null @@ -1,240 +0,0 @@ -import os -import time -import pickle -from multiprocessing import Process - -from mmcv import build_from_cfg - -from mpa.registry import STAGES -from mpa.stage import Stage -from mpa.utils.logger import get_logger - -logger = get_logger() - - -def build_train_stage(trainer_stage_config, common_cfg): - if trainer_stage_config is None: - raise ValueError( - 'HpoRunner needs a trainer config to run.' - ) - - if trainer_stage_config.get('type', None) is None: - raise ValueError( - 'The type of hpo trainer should be specified.' - ) - if trainer_stage_config.get('config', None) is None: - raise ValueError( - 'The file path of hpo trainer config should be specified.' - ) - trainer_stage_config.name = 'hpo_trainer' - trainer_stage_config.mode = 'train' - trainer_stage_config.common_cfg = common_cfg - - return build_from_cfg(trainer_stage_config, STAGES) - - -def run_hpo_trainer(trainer_stage_config, common_cfg, metric, mode, - model_cfg, data_cfg, model_ckpt, - output_path, hp_config): - trainer_stage = build_train_stage(trainer_stage_config, common_cfg) - - # replace old hyper-parameters with the newly created ones. - if 'lr' in hp_config['params']: - trainer_stage.cfg.optimizer.lr = hp_config['params']['lr'] - - if 'bs' in hp_config['params']: - data_cfg.data.samples_per_gpu = int(hp_config['params']['bs']) - - if hasattr(trainer_stage.cfg, 'runner'): - trainer_stage.cfg.runner.max_epochs = hp_config['iterations'] - trainer_stage.cfg.total_epochs = hp_config['iterations'] - - # add HPOHook - hpo_hook = dict(type='HPOHook', - hp_config=hp_config, - metric=metric, - priority='LOW') - custom_hooks = trainer_stage.cfg.get('custom_hooks', []) - custom_hooks.append(hpo_hook) - trainer_stage.cfg['custom_hooks'] = custom_hooks - - _ = trainer_stage.run(stage_idx=hp_config['trial_id'], - mode=mode, - model_cfg=model_cfg, - data_cfg=data_cfg, - model_ckpt=model_ckpt, - output_path=output_path, - hp_config=hp_config) - - -def exec_hpo_trainer(arg_file_name, gpu_id): - trainer_file_name = os.path.join(os.path.dirname(__file__), 'hpo_trainer.py') - os.system(f'CUDA_VISIBLE_DEVICES={gpu_id} PYTHONPATH="{os.path.dirname(__file__)}/../../:$PYTHONPATH" ' - f'python {trainer_file_name} {arg_file_name}') - time.sleep(10) - return gpu_id - - -@STAGES.register_module() -class HpoRunner(Stage): - def __init__(self, **kwargs): - if kwargs.get('config', None) is None: - kwargs['config'] = {'hpo': {'trainer': {'type': None, 'config': None}}} - - super().__init__(**kwargs) - self.cfg['hpo'] = kwargs.pop('hpo', None) - self.cfg['common_cfg'] = kwargs.pop('common_cfg', None) - - if self.cfg['hpo'] is None: - raise ValueError( - 'Missing HPO configuration.' - ) - - self.hyperparams = self.cfg.hpo.get('hyperparams', None) - if self.hyperparams is None: - raise ValueError( - 'The list of hyper-parameters to tune is required to run HPO.' - ) - - self.metric = self.cfg.hpo.get('metric', None) - if self.metric is None: - raise ValueError( - 'The name of metric is required to run HPO.' - ) - self.num_trials = self.cfg.hpo.get('num_trials', 10) - self.max_iterations = self.cfg.hpo.get('max_iterations', 10) - self.subset_size = self.cfg.hpo.get('subset_size', 0) - self.search_alg = self.cfg.hpo.get('search_alg', 'bayes_opt') - if self.search_alg not in ['bayes_opt', 'asha']: - raise ValueError( - 'The \'search_alg\' should be one of \'bayes_opt\' or \'asha\'.' - ) - self.gpu_list = self.cfg.hpo.get('gpu_list', [0]) - self.num_workers = len(self.gpu_list) - self.num_brackets = self.cfg.hpo.get('num_brackets', 1) - self.min_iterations = self.cfg.hpo.get('min_iterations', 1) - self.reduction_factor = self.cfg.hpo.get('reduction_factor', 4) - - def run(self, model_cfg, model_ckpt, data_cfg, **kwargs): - import hpopt - - mode = kwargs.get('mode', 'train') - if mode not in self.mode: - return {} - - output_path = self.cfg.work_dir - - hparams = {} - for param in self.hyperparams: - hparams[param.name] = hpopt.search_space(param.type, param.range) - - hpoptimizer = None - - if self.search_alg == 'bayes_opt': - hpoptimizer = hpopt.create(save_path=output_path, - search_alg='bayes_opt', - search_space=hparams, - ealry_stop='median_stop', - num_init_trials=5, - num_trials=self.num_trials, - max_iterations=self.max_iterations, - subset_size=self.subset_size) - elif self.search_alg == 'asha': - hpoptimizer = hpopt.create(save_path=output_path, - search_alg='asha', - search_space=hparams, - num_trials=self.num_trials, - min_iterations=self.min_iterations, - max_iterations=self.max_iterations, - reduction_factor=self.reduction_factor, - num_brackets=self.num_brackets, - subset_size=self.subset_size) - - logger.info('** HPO START **') - - if self.search_alg == 'bayes_opt': - while True: - hp_config = hpoptimizer.get_next_sample() - - if hp_config is None: - break - - run_hpo_trainer(trainer_stage_config=self.cfg.hpo.trainer.copy(), - common_cfg=self.cfg['common_cfg'].copy(), - metric=self.metric, - mode=mode, - model_cfg=model_cfg, - data_cfg=data_cfg, - model_ckpt=model_ckpt, - output_path=output_path, - hp_config=hp_config) - elif self.search_alg == 'asha': - proc_list = [] - gpu_alloc_list = [] - - while True: - num_active_workers = 0 - for p, g in zip(reversed(proc_list), reversed(gpu_alloc_list)): - if p.is_alive(): - num_active_workers += 1 - else: - p.close() - proc_list.remove(p) - gpu_alloc_list.remove(g) - - if num_active_workers == self.num_workers: - time.sleep(10) - - while num_active_workers < self.num_workers: - hp_config = hpoptimizer.get_next_sample() - - if hp_config is None: - break - - _kwargs = {"trainer_stage_config": self.cfg.hpo.trainer, - "common_cfg": self.cfg['common_cfg'], - "metric": self.metric, - "mode": mode, - "model_cfg": model_cfg, - "data_cfg": data_cfg, - "model_ckpt": model_ckpt, - "output_path": output_path, - "hp_config": hp_config} - - pickle_path = os.path.join(output_path, f"hpo_trial_{hp_config['trial_id']}.pickle") - with open(pickle_path, "wb") as pfile: - pickle.dump(_kwargs, pfile) - - gpu_id = -1 - for idx in self.gpu_list: - if idx not in gpu_alloc_list: - gpu_id = idx - break - - if gpu_id < 0: - raise ValueError('No available GPU!!') - - p = Process(target=exec_hpo_trainer, args=(pickle_path, gpu_id, )) - proc_list.append(p) - gpu_alloc_list.append(gpu_id) - p.start() - num_active_workers += 1 - - # All trials are done. - if num_active_workers == 0: - break - - best_config = hpoptimizer.get_best_config() - print("best hp: ", best_config) - - logger.info('** HPO END **') - - retval = {'hyperparams': {}} - - if 'lr' in best_config: - retval['hyperparams']['lr'] = best_config.get('lr') - - if 'bs' in best_config: - retval['hyperparams']['bs'] = best_config.get('bs') - - return retval diff --git a/mpa/utils/hpo_trainer.py b/mpa/utils/hpo_trainer.py deleted file mode 100644 index 68f7b885..00000000 --- a/mpa/utils/hpo_trainer.py +++ /dev/null @@ -1,8 +0,0 @@ -from mpa.utils.hpo_stage import run_hpo_trainer -import pickle -import sys - -if __name__ == '__main__': - with open(sys.argv[1], "rb") as pfile: - kwargs = pickle.load(pfile) - run_hpo_trainer(**kwargs) diff --git a/mpa/utils/mda_stage.py b/mpa/utils/mda_stage.py deleted file mode 100644 index 6f8af3e8..00000000 --- a/mpa/utils/mda_stage.py +++ /dev/null @@ -1,178 +0,0 @@ -import os.path as osp - -from torchvision.transforms import ToPILImage, ToTensor, Resize - -from mmcv.parallel import MMDataParallel -from mmcv.runner import wrap_fp16_model -from mmcls.datasets import build_dataloader as build_dataloader_cls -from mmcls.datasets import build_dataset as build_dataset_cls -from mmcls.models import build_classifier -from mmcls.datasets import PIPELINES as PIPELINES_mmcls -from mmdet.datasets import build_dataloader as build_dataloader_det -from mmdet.datasets import build_dataset as build_dataset_det -from mmdet.datasets import PIPELINES as PIPELINES_mmdet -from mmdet.models import build_detector - -from mda import MDA - -from mpa.cls.stage import ClsStage -from mpa.det.stage import DetectionStage -from mpa.det.inferrer import replace_ImageToTensor -from mpa.registry import STAGES -from mpa.stage import Stage -from mpa.utils.logger import get_logger - -logger = get_logger() - -TASK_STAGE = {'classification': ClsStage, 'detection': DetectionStage} - - -def fix_dataset_size(results): - to_tensor = ToTensor() - to_pil = ToPILImage() - resize = Resize([128, 128]) - results = to_tensor(resize(to_pil(results))) - - return results - - -@PIPELINES_mmcls.register_module() -class ClsFixDatasetSize(object): - def __call__(self, results): - results['img'] = fix_dataset_size(results['img']) - return results - - -@PIPELINES_mmdet.register_module() -class DetFixDatasetSize(object): - def __call__(self, results): - results['img'][0] = fix_dataset_size(results['img'][0]) - return results - - -@STAGES.register_module() -class MdaRunner(Stage): - def __init__(self, **kwargs): - super().__init__(**kwargs) - self.task = kwargs.get('task', None) # classification or detection - if self.task not in ['classification', 'detection']: - raise ValueError(f'task shoulde be "classification" or "detection". Current value is {self.task}') - self.task_stage = TASK_STAGE[self.task](**kwargs) - - mda_metric = kwargs.get('mda_metric', 'z-score') # 'z-score', 'cos-sim', 'kl' or 'wst' - if mda_metric not in ['z-score', 'cos-sim', 'kl', 'wst']: - raise ValueError(f'mda_metric should be "z-score", "cos-sim", "kl" or "wst". \ - Current value is "{mda_metric}"') - self.mda_metric = mda_metric - - def make_dataset(self, cfg): - if self.task == "classification": - input_source = cfg.get('input_source', 'test') - print(f'MDA on input source: data.{input_source}') - cfg.data[input_source]['pipeline'].append({'type': 'ClsFixDatasetSize'}) - dataset = build_dataset_cls(cfg.data[input_source]) - elif self.task == "detection": - input_source = cfg.get('input_source', 'test') - # Current version of OTE detection supports only batch size 1 during evaluation and inference. - # When version is updated, this code will be fixed. - samples_per_gpu = 1 - - # Input source - input_source = cfg.get('input_source', 'test') - print(f'MDA on input source: data.{input_source}') - if input_source == 'train': - src_data_cfg = DetectionStage.get_train_data_cfg(cfg) - else: - src_data_cfg = cfg.data[input_source] - - if samples_per_gpu > 1: # Replace 'ImageToTensor' to 'DefaultFormatBundle' - src_data_cfg.pipeline = replace_ImageToTensor(src_data_cfg.pipeline) - - data_cfg = src_data_cfg.copy() - dataset = build_dataset_det(data_cfg) - - return dataset - - def make_model(self, cfg): - if self.task == "classification": - # build the model and load checkpoint - model = build_classifier(cfg.model) - elif self.task == "detection": - # Target classes - if 'task_adapt' in cfg: - target_classes = cfg.task_adapt.final - else: - target_classes = self.dataset.CLASSES - - # Model - cfg.model.pretrained = None - if cfg.model.get('neck'): - if isinstance(cfg.model.neck, list): - for neck_cfg in cfg.model.neck: - if neck_cfg.get('rfp_backbone'): - if neck_cfg.rfp_backbone.get('pretrained'): - neck_cfg.rfp_backbone.pretrained = None - elif cfg.model.neck.get('rfp_backbone'): - if cfg.model.neck.rfp_backbone.get('pretrained'): - cfg.model.neck.rfp_backbone.pretrained = None - - train_cfg = cfg.get('train_cfg', None) - test_cfg = cfg.get('test_cfg', None) - model = build_detector(cfg.model, train_cfg=train_cfg, test_cfg=test_cfg) - - model.CLASSES = target_classes - - fp16_cfg = cfg.get('fp16', None) - if fp16_cfg is not None: - wrap_fp16_model(model) - return model - - def analyze_model_drift(self, cfg): - self.dataset = self.make_dataset(cfg) - - # Data loader - if self.task == "classification": - data_loader = build_dataloader_cls( - self.dataset, - samples_per_gpu=cfg.data.samples_per_gpu, - workers_per_gpu=cfg.data.workers_per_gpu, - dist=False, - shuffle=False, - round_up=False) - elif self.task == "detection": - data_loader = build_dataloader_det( - self.dataset, - samples_per_gpu=cfg.data.test.get('samples_per_gpu', 1), - workers_per_gpu=cfg.data.workers_per_gpu, - dist=False, - shuffle=False) - - self.model = self.make_model(cfg) - self.model = MMDataParallel(self.model, device_ids=[0]) - mda = MDA(self.mda_metric, verbose=1, mode='stable') - output = mda.measure(self.model, data_loader) - - print(f'Model drift score is {output}') - with open(osp.join(osp.abspath(cfg.work_dir), 'mpa_output.txt'), 'w') as f: - f.write(str(output)) - - return output - - def run(self, model_cfg, model_ckpt, data_cfg, **kwargs): - """Run evaluation stage - - - Run inference - - Run evaluation via MMDetection -> MMCV - """ - mode = kwargs.get('mode', 'train') - if mode not in self.mode: - logger.warning(f'mode for this stage {mode}') - return {} - - self.task_stage._init_logger() - cfg = self.task_stage.configure(model_cfg, model_ckpt, data_cfg, training=False, **kwargs) - - # mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir)) - mda_results = self.analyze_model_drift(cfg) - - return mda_results From 80432c86e42b6c3a83a662881a390bf7c6a83fa4 Mon Sep 17 00:00:00 2001 From: Songki Choi Date: Fri, 18 Mar 2022 00:05:51 +0900 Subject: [PATCH 2/2] Remove OMZ/Experimental features from external release --- mpa/modules/__init__.py | 2 - mpa/modules/experimental/__init__.py | 6 - mpa/modules/experimental/datasets/__init__.py | 8 - .../adaptive_equalization_sampling_dataset.py | 53 -- .../experimental/datasets/cls_tv_dataset.py | 264 ---------- .../datasets/oversampling_dataset.py | 158 ------ .../datasets/pipelines/__init__.py | 2 - .../datasets/pipelines/transforms/__init__.py | 2 - .../pipelines/transforms/random_args_segm.py | 174 ------- .../datasets/pseudo_incr_dataset.py | 254 ---------- .../datasets/pseudo_semi_dataset.py | 63 --- mpa/modules/experimental/hooks/__init__.py | 3 - .../hooks/confidence_bank_hook.py | 46 -- mpa/modules/experimental/hooks/stop_hook.py | 12 - mpa/modules/experimental/models/__init__.py | 10 - .../experimental/models/detectors/__init__.py | 3 - .../detectors/custom_retinanet_detector.py | 11 - .../detectors/lwf_two_stage_detector.py | 87 ---- .../experimental/models/heads/__init__.py | 6 - .../models/heads/lwf_bbox_head.py | 449 ----------------- .../experimental/models/heads/lwf_roi_head.py | 232 --------- .../experimental/models/heads/lwf_rpn_head.py | 89 ---- .../models/heads/pseudo_bbox_head.py | 113 ----- .../models/heads/reco_sep_aspp_head.py | 187 ------- .../experimental/models/losses/__init__.py | 11 - .../experimental/models/losses/dice_loss.py | 165 ------ .../models/losses/distillation_losses.py | 55 -- .../experimental/models/losses/dmt_loss.py | 229 --------- .../experimental/models/losses/eqlv2.py | 96 ---- .../experimental/models/losses/focal_loss.py | 54 -- .../experimental/models/losses/kd_loss.py | 67 --- .../experimental/models/losses/lwf_loss.py | 191 ------- .../experimental/models/losses/reco_loss.py | 141 ------ .../models/losses/reverse_loss.py | 193 ------- .../models/losses/softiou_loss.py | 140 ------ .../models/segmentors/__init__.py | 7 - .../models/segmentors/cutmix/gt_utils.py | 168 ------- .../models/segmentors/cutmix/mask_gen.py | 74 --- .../models/segmentors/cutmix_seg_gt_naive.py | 129 ----- .../models/segmentors/cutmix_seg_naive.py | 115 ----- .../models/segmentors/keep_output.py | 77 --- .../models/segmentors/mean_teacher_naive.py | 93 ---- .../models/segmentors/oversampling.py | 29 -- .../models/segmentors/semisl_segmentor.py | 120 ----- mpa/modules/experimental/utils/pseudo_nms.py | 63 --- mpa/modules/omz/__init__.py | 2 - mpa/modules/omz/custom_layers/__init__.py | 0 mpa/modules/omz/custom_layers/base_layers.py | 369 -------------- mpa/modules/omz/models/__init__.py | 4 - mpa/modules/omz/models/anchor/__init__.py | 1 - .../omz/models/anchor/anchor_generator.py | 47 -- mpa/modules/omz/models/backbones/__init__.py | 3 - .../omz/models/backbones/omz_backbone.py | 74 --- .../omz/models/backbones/omz_backbone_cls.py | 49 -- .../omz/models/backbones/omz_backbone_det.py | 57 --- .../omz/models/dense_heads/__init__.py | 3 - .../omz/models/dense_heads/omz_rpn_head.py | 57 --- .../omz/models/dense_heads/omz_ssd_head.py | 73 --- mpa/modules/omz/utils/__init__.py | 0 mpa/modules/omz/utils/omz_to_torch.py | 473 ------------------ mpa/modules/omz/utils/omz_utils.py | 90 ---- mpa/modules/omz/utils/utils.py | 83 --- 62 files changed, 5836 deletions(-) delete mode 100644 mpa/modules/experimental/__init__.py delete mode 100644 mpa/modules/experimental/datasets/__init__.py delete mode 100644 mpa/modules/experimental/datasets/adaptive_equalization_sampling_dataset.py delete mode 100644 mpa/modules/experimental/datasets/cls_tv_dataset.py delete mode 100644 mpa/modules/experimental/datasets/oversampling_dataset.py delete mode 100644 mpa/modules/experimental/datasets/pipelines/__init__.py delete mode 100644 mpa/modules/experimental/datasets/pipelines/transforms/__init__.py delete mode 100644 mpa/modules/experimental/datasets/pipelines/transforms/random_args_segm.py delete mode 100644 mpa/modules/experimental/datasets/pseudo_incr_dataset.py delete mode 100644 mpa/modules/experimental/datasets/pseudo_semi_dataset.py delete mode 100644 mpa/modules/experimental/hooks/__init__.py delete mode 100644 mpa/modules/experimental/hooks/confidence_bank_hook.py delete mode 100644 mpa/modules/experimental/hooks/stop_hook.py delete mode 100644 mpa/modules/experimental/models/__init__.py delete mode 100644 mpa/modules/experimental/models/detectors/__init__.py delete mode 100644 mpa/modules/experimental/models/detectors/custom_retinanet_detector.py delete mode 100644 mpa/modules/experimental/models/detectors/lwf_two_stage_detector.py delete mode 100644 mpa/modules/experimental/models/heads/__init__.py delete mode 100644 mpa/modules/experimental/models/heads/lwf_bbox_head.py delete mode 100644 mpa/modules/experimental/models/heads/lwf_roi_head.py delete mode 100644 mpa/modules/experimental/models/heads/lwf_rpn_head.py delete mode 100644 mpa/modules/experimental/models/heads/pseudo_bbox_head.py delete mode 100644 mpa/modules/experimental/models/heads/reco_sep_aspp_head.py delete mode 100644 mpa/modules/experimental/models/losses/__init__.py delete mode 100644 mpa/modules/experimental/models/losses/dice_loss.py delete mode 100644 mpa/modules/experimental/models/losses/distillation_losses.py delete mode 100644 mpa/modules/experimental/models/losses/dmt_loss.py delete mode 100644 mpa/modules/experimental/models/losses/eqlv2.py delete mode 100644 mpa/modules/experimental/models/losses/focal_loss.py delete mode 100644 mpa/modules/experimental/models/losses/kd_loss.py delete mode 100644 mpa/modules/experimental/models/losses/lwf_loss.py delete mode 100644 mpa/modules/experimental/models/losses/reco_loss.py delete mode 100644 mpa/modules/experimental/models/losses/reverse_loss.py delete mode 100644 mpa/modules/experimental/models/losses/softiou_loss.py delete mode 100644 mpa/modules/experimental/models/segmentors/__init__.py delete mode 100644 mpa/modules/experimental/models/segmentors/cutmix/gt_utils.py delete mode 100644 mpa/modules/experimental/models/segmentors/cutmix/mask_gen.py delete mode 100644 mpa/modules/experimental/models/segmentors/cutmix_seg_gt_naive.py delete mode 100644 mpa/modules/experimental/models/segmentors/cutmix_seg_naive.py delete mode 100644 mpa/modules/experimental/models/segmentors/keep_output.py delete mode 100644 mpa/modules/experimental/models/segmentors/mean_teacher_naive.py delete mode 100644 mpa/modules/experimental/models/segmentors/oversampling.py delete mode 100644 mpa/modules/experimental/models/segmentors/semisl_segmentor.py delete mode 100644 mpa/modules/experimental/utils/pseudo_nms.py delete mode 100644 mpa/modules/omz/__init__.py delete mode 100644 mpa/modules/omz/custom_layers/__init__.py delete mode 100644 mpa/modules/omz/custom_layers/base_layers.py delete mode 100644 mpa/modules/omz/models/__init__.py delete mode 100644 mpa/modules/omz/models/anchor/__init__.py delete mode 100644 mpa/modules/omz/models/anchor/anchor_generator.py delete mode 100644 mpa/modules/omz/models/backbones/__init__.py delete mode 100644 mpa/modules/omz/models/backbones/omz_backbone.py delete mode 100644 mpa/modules/omz/models/backbones/omz_backbone_cls.py delete mode 100644 mpa/modules/omz/models/backbones/omz_backbone_det.py delete mode 100644 mpa/modules/omz/models/dense_heads/__init__.py delete mode 100644 mpa/modules/omz/models/dense_heads/omz_rpn_head.py delete mode 100644 mpa/modules/omz/models/dense_heads/omz_ssd_head.py delete mode 100644 mpa/modules/omz/utils/__init__.py delete mode 100644 mpa/modules/omz/utils/omz_to_torch.py delete mode 100644 mpa/modules/omz/utils/omz_utils.py delete mode 100644 mpa/modules/omz/utils/utils.py diff --git a/mpa/modules/__init__.py b/mpa/modules/__init__.py index 095f8e9c..64f600e0 100644 --- a/mpa/modules/__init__.py +++ b/mpa/modules/__init__.py @@ -2,6 +2,4 @@ from . import datasets from . import hooks from . import models -from . import omz from . import optimizer -from . import experimental diff --git a/mpa/modules/experimental/__init__.py b/mpa/modules/experimental/__init__.py deleted file mode 100644 index 8bb791b5..00000000 --- a/mpa/modules/experimental/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# flake8: noqa -from . import datasets -from . import hooks -from . import models -# from . import omz -# from . import optimizer diff --git a/mpa/modules/experimental/datasets/__init__.py b/mpa/modules/experimental/datasets/__init__.py deleted file mode 100644 index 0559465a..00000000 --- a/mpa/modules/experimental/datasets/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -# flake8: noqa -from . import cls_tv_dataset -from . import pseudo_incr_dataset -from . import pseudo_semi_dataset -from . import oversampling_dataset -from . import adaptive_equalization_sampling_dataset -from . import cls_tv_dataset -from . import pipelines \ No newline at end of file diff --git a/mpa/modules/experimental/datasets/adaptive_equalization_sampling_dataset.py b/mpa/modules/experimental/datasets/adaptive_equalization_sampling_dataset.py deleted file mode 100644 index a5b94df5..00000000 --- a/mpa/modules/experimental/datasets/adaptive_equalization_sampling_dataset.py +++ /dev/null @@ -1,53 +0,0 @@ -import torch -import torch.multiprocessing as mp - -from mmseg.datasets import DATASETS, build_dataset - - -@DATASETS.register_module() -class AdaptiveEqualizationSamplingDataset(object): - """Dataset wrapper for Adaptive Equalization Sampling. - """ - - def __init__(self, - orig_type=None, - **kwargs): - - dataset_cfg = kwargs.copy() - dataset_cfg['type'] = orig_type - self.dataset = build_dataset(dataset_cfg) - self.num_samples = len(self.dataset) - - # why is it used with (3, class) shape? - self._class_criterion = torch.rand(len(self.dataset.CLASSES)).type(torch.float32) - print('####################### AdaptiveEqualizationSamplingDataset') - - def update_confidence_bank(self, conf): - print(id(self), 'before :', self._class_criterion) - self._class_criterion = conf - # dist.broadcast(self._class_criterion, 0) - print(id(self), 'after :', self._class_criterion) - print(dir(mp)) - print(mp.active_children()) - - ctx = mp.get_context('spawn') - print(ctx) - for children in mp.active_children(): - print(children) - print(dir(children)) - - @property - def class_criterion(self): - return self._class_criterion - - def __len__(self): - """Total number of samples of data.""" - return self.num_samples - - def __getitem__(self, idx): - # sample_from_bank - worker_info = torch.utils.data.get_worker_info() - - print(worker_info, 'self._class_criterion :', self._class_criterion) - print(worker_info.dataset.class_criterion) - return self.dataset.__getitem__(idx) diff --git a/mpa/modules/experimental/datasets/cls_tv_dataset.py b/mpa/modules/experimental/datasets/cls_tv_dataset.py deleted file mode 100644 index 628e07a7..00000000 --- a/mpa/modules/experimental/datasets/cls_tv_dataset.py +++ /dev/null @@ -1,264 +0,0 @@ -import os -import copy -import pickle -import numpy as np - -import torchvision - -from mmcv.utils.registry import build_from_cfg -from mmcls.datasets.builder import DATASETS, PIPELINES -from mmcls.datasets.base_dataset import BaseDataset -from mmcls.datasets.pipelines import Compose - -from mpa.modules.datasets.sources.index_list import ImageIndexList -from mpa.utils.logger import get_logger - -logger = get_logger() - - -@DATASETS.register_module() -class ClsTVDataset(BaseDataset): - """ create a dataset as a subset from torchvision dataset using - include and/or exclude image indices (list or pickle file). - """ - def __init__(self, base, classes=[], new_classes=[], data_prefix=None, **kwargs): - self.num_images = kwargs.pop('num_images', -1) - self._samples_per_gpu = kwargs.pop('samples_per_gpu', 1) - self._workers_per_gpu = kwargs.pop('workers_per_gpu', 1) - pipeline = kwargs.pop('pipeline', None) - include_idx = kwargs.pop('include_idx', None) - exclude_idx = kwargs.pop('exclude_idx', None) - self.use_labels = kwargs.pop('use_labels', True) - self.balanced_class = kwargs.pop('balanced_class', True) - _ = kwargs.pop('test_mode', False) - self.dataset_type = base - self.class_acc = False - self.img_indices = dict(old=[], new=[]) - - if 'download' not in kwargs: - kwargs['download'] = True - - # Get Dataset from torchvision.datasets - if isinstance(base, str): - if not data_prefix: - data_prefix = os.path.join('data/torchvision', base.lower()) - self.base = getattr(torchvision.datasets, base)(data_prefix, **kwargs) - else: - self.base = base - - # Get Labels from torchvision.datasets - if hasattr(self.base, 'targets'): - self.labels = np.array(self.base.targets) - elif hasattr(self.base, 'labels'): - self.labels = np.array(self.base.labels) - else: - raise NotImplementedError(f'check the attribute name of torchvision \ - dataset {base}') - if not classes: - classes = list(set(self.labels.tolist())) - self.CLASSES = classes - self.new_classes = new_classes - - # Configuration indices pool - indices_pool = range(len(self.labels)) - indices_pool, self.exclude_idx = self.configure_idx(indices_pool, include_idx, exclude_idx) - - # Sampling indices - self.indices = self.sampling_idx(indices_pool) - self.data_source = ImageIndexList(self.base, self.indices) - - # configuration Pipelines - self.pipeline, self.num_pipes = self.configure_pipeline(pipeline) - - # Load Annotations - self.data_infos = self.load_annotations() - self.statistics() - - def statistics(self): - logger.info(f'ClsTVDataset - {len(self.CLASSES)} classes from {self.dataset_type}') - logger.info(f'- Classes: {self.CLASSES}') - if self.new_classes: - logger.info(f'- New Classes: {self.new_classes}') - old_data_length = len(self.img_indices['old']) - new_data_length = len(self.img_indices['new']) - logger.info(f'- # of old classes images: {old_data_length}') - logger.info(f'- # of New classes images: {new_data_length}') - logger.info(f'- # of images: {len(self)}') - - @staticmethod - def configure_idx(indices_pool, include_idx, exclude_idx): - if include_idx is not None: - if isinstance(include_idx, list): - indices_pool = include_idx - elif isinstance(include_idx, str): - if os.path.exists(include_idx): - indices_pool = pickle.load(open(include_idx, 'rb')) - else: - logger.warning(f'cannot find include index pickle file \ - {include_idx}. ignored.') - else: - raise TypeError(f"not supported type for 'include_idx'.\ - should be list or pickle file path but {type(include_idx)}") - - if exclude_idx is not None: - if isinstance(exclude_idx, str): - if os.path.exists(exclude_idx): - exclude_idx = pickle.load(open(exclude_idx, 'rb')) - else: - logger.warning(f'cannot find exclude index pickle file \ - {exclude_idx}. ignored.') - elif not isinstance(exclude_idx, list): - raise TypeError(f"not supported type for 'exclude_idx'.\ - should be list or pickle file path but {type(exclude_idx)}") - if isinstance(exclude_idx, list): - indices_pool = np.setdiff1d(indices_pool, exclude_idx) - return indices_pool, exclude_idx - - def sampling_idx(self, indices_pool): - indices = [] - if self.num_images != -1: - if self.num_images > len(indices_pool) or self.num_images < 0: - raise RuntimeError(f"cannot generate split dataset. \ - length of base dataset = {len(indices_pool)}, \ - requested split {self.num_images}") - # if self.balanced_class: - items_per_class = self.num_images // len(self.CLASSES) - for i in self.CLASSES: - idx = np.where(self.labels == i)[0] - idx = list(set(idx) & set(indices_pool)) - if self.balanced_class: - idx = np.random.choice(idx, items_per_class, False) - np.random.shuffle(idx) - indices.extend(idx) - indices = np.array(indices) - if not self.balanced_class: - indices = np.random.choice(indices, self.num_images, False) - else: - for i in self.CLASSES: - idx = np.where(self.labels == i)[0] - idx = list(set(idx) & set(indices_pool)) - np.random.shuffle(idx) - indices.extend(idx) - indices = np.array(indices) - return indices - - @staticmethod - def configure_pipeline(pipeline_cfg): - num_pipes = 1 - pipeline = {} - if pipeline_cfg is None: - # set default pipeline - pipeline_cfg = [ - dict(type='Normalize', mean=[127., 127., 127.], std=[127., 127., 127.]), - dict(type='ImageToTensor', keys=['img']), - dict(type='ToTensor', keys=['gt_label']), - dict(type='Collect', keys=['img', 'gt_label']) - ] - - if isinstance(pipeline_cfg, dict): - for k, v in pipeline_cfg.items(): - pipeline[k] = \ - Compose([build_from_cfg(p, PIPELINES) for p in v]) - logger.debug(pipeline[k]) - num_pipes = len(pipeline_cfg) - elif isinstance(pipeline_cfg, list): - if len(pipeline_cfg) <= 0: - pipeline = None - else: - pipeline = \ - Compose([build_from_cfg(p, PIPELINES) for p in pipeline_cfg]) - return pipeline, num_pipes - - def load_annotations(self): - data_infos = [] - for idx in range(self.data_source.get_length()): - img, label = self.data_source.get_sample(idx) - gt_label = self.class_to_idx[label] - info = {'img': img, 'gt_label': np.array(gt_label, dtype=np.int64)} - data_infos.append(info) - if label in self.new_classes: - self.img_indices['new'].append(idx) - else: - self.img_indices['old'].append(idx) - return data_infos - - def __len__(self): - return len(self.data_infos) - - def __getitem__(self, idx): - if self.pipeline is None: - return self.data_infos[idx] - - data_infos = [ - copy.deepcopy(self.data_infos[idx]) for _ in range(self.num_pipes) - ] - if isinstance(self.pipeline, dict): - results = {} - for i, (k, v) in enumerate(self.pipeline.items()): - results[k] = self.pipeline[k](data_infos[i]) - else: - results = self.pipeline(data_infos[0]) - - return results - - def evaluate(self, - results, - metric='accuracy', - metric_options=None, - logger=None): - """Evaluate the dataset with new metric 'class_accuracy' - - Args: - results (list): Testing results of the dataset. - metric (str | list[str]): Metrics to be evaluated. - Default value is `accuracy`. - 'accuracy', 'precision', 'recall', 'f1_score', 'support', 'class_accuracy' - metric_options (dict, optional): Options for calculating metrics. - Allowed keys are 'topk', 'thrs' and 'average_mode'. - Defaults to None. - logger (logging.Logger | str, optional): Logger used for printing - related information during evaluation. Defaults to None. - Returns: - dict: evaluation results - """ - if metric_options is None: - metric_options = {'topk': (1, 5) if len(self.CLASSES) >= 5 else (1, )} - - if isinstance(metric, str): - metrics = [metric] - else: - metrics = metric - - if 'class_accuracy' in metrics: - metrics.remove('class_accuracy') - self.class_acc = True - - eval_results = super().evaluate(results, metrics, metric_options, logger) - - # Add Evaluation Accuracy score per Class - if self.class_acc: - results = np.vstack(results) - gt_labels = self.get_gt_labels() - accuracies = self.class_accuracy(results, gt_labels) - eval_results.update({f'{c} accuracy': a for c, a in zip(self.CLASSES, accuracies)}) - eval_results.update({'mean accuracy': np.mean(accuracies)}) - - return eval_results - - def class_accuracy(self, results, gt_labels): - accracies = [] - pred_label = results.argsort(axis=1)[:, -1:][:, ::-1] - for i in range(len(self.CLASSES)): - cls_pred = pred_label == i - cls_pred = cls_pred[gt_labels == i] - cls_acc = np.sum(cls_pred) / len(cls_pred) - accracies.append(cls_acc) - return accracies - - @property - def samples_per_gpu(self): - return self._samples_per_gpu - - @property - def workers_per_gpu(self): - return self._workers_per_gpu diff --git a/mpa/modules/experimental/datasets/oversampling_dataset.py b/mpa/modules/experimental/datasets/oversampling_dataset.py deleted file mode 100644 index 74c0d5cc..00000000 --- a/mpa/modules/experimental/datasets/oversampling_dataset.py +++ /dev/null @@ -1,158 +0,0 @@ -import os -import cv2 -import tqdm -import numpy as np -import pandas as pd - -import torch -from mmseg.datasets import DATASETS, build_dataset -from mpa.utils.logger import get_logger - -logger = get_logger() - - -@DATASETS.register_module() -class OversamplingSampler(object): - """Dataset wrapper for oversampling for Imbalanced dataset. - """ - - def __init__(self, - orig_type=None, - sampling_mode='whole', - weights_mode='max', - weights_file='oversampling_weights.csv', - **kwargs): - - assert orig_type is not None - assert sampling_mode in ['whole', 'half', 'strict_half'] - assert weights_mode in ['max', 'sum', 'adaptive'] - - if sampling_mode == 'strict_half': - logger.info("In 'strict_half' sampling mode, OverSamplingSegmentor and half batch_size are required.") - logger.info("If samples_per_gpu is set to 8, REAL samples_per_gpu will be 16.") - - self.sampling_mode = sampling_mode - - dataset_cfg = kwargs.copy() - dataset_cfg['type'] = orig_type - self.dataset = build_dataset(dataset_cfg) - self.extension = self.dataset.img_infos[0]['filename'].split('.')[-1] - - if os.path.isfile(os.path.join(kwargs['data_root'], weights_file)): - self.df_dataset = pd.read_csv(os.path.join(kwargs['data_root'], weights_file), - index_col='Unnamed: 0') - - if not [self.dataset.img_infos[i]['filename'] - for i in range(len(self.dataset.img_infos))] == self.df_dataset.index.tolist(): - - logger.info(("The order of self.dataset and the previously saved order of " - "self.df_dataset do not match.")) - logger.info("self.df_dataset is reset.") - self.set_df() - self.df_dataset.to_csv(os.path.join(kwargs['data_root'], weights_file)) - logger.info("self.df_dataset is saved at {}.".format( - os.path.join(kwargs['data_root'], weights_file))) - - else: - logger.info("self.df_dataset is set.") - self.set_df() - self.df_dataset.to_csv(os.path.join(kwargs['data_root'], weights_file)) - logger.info("self.df_dataset is saved at {}.".format( - os.path.join(kwargs['data_root'], weights_file))) - - self.num_samples = len(self.dataset) - - self.init_weights(weights_mode=weights_mode) - self.init_counts() - - def set_df(self): - # ann_list = os.listdir(self.dataset.ann_dir) - # ann_extension = ann_list[0].split('.')[-1] - ann_infos = {} - for path in tqdm.tqdm(self.dataset.img_infos, desc='Get # of classes', total=len(self.dataset.img_infos)): - img_path, ann_path = path['filename'], path['ann']['seg_map'] - if img_path not in ann_infos: - # img_path = path.replace(ann_extension, self.extension) - ann_infos[img_path] = \ - [0 for _ in range(len(self.dataset.CLASSES))] - mask = cv2.imread(os.path.join(self.dataset.ann_dir, ann_path)) - labels = np.unique(mask) - for label in labels: - if label != 255: - ann_infos[img_path][label] = 1 - - # ann_infos = dict(sorted(ann_infos.items(), key=lambda x: x[0])) - self.df_dataset = pd.DataFrame(ann_infos).T - - def init_weights(self, weights_mode='max'): - weights = (1 / self.df_dataset.sum(axis=0)).tolist() - if weights_mode == 'max': - self.df_dataset['weights'] = 0 - for i in range(len(self.dataset.CLASSES)): - self.df_dataset['weights'] = [ - max(*elem) for elem in zip( - self.df_dataset['weights'], - self.df_dataset.iloc[:, i].apply(lambda x: weights[i] if x else 0) - ) - ] - - elif weights_mode == 'sum': - raise NotImplementedError() - # for i in range(self.num_samples): - # if i == 0: - # result = self.df_dataset.loc[:,i].apply( - # lambda x: weights[i] if x else 0) - # else: - # result += self.df_dataset.loc[:,i].apply( - # lambda x: weights[i] if x else 0) - # self.df_dataset['weights'] = result - - elif weights_mode == 'adaptive': - raise NotImplementedError() - - self.weights = torch.DoubleTensor(self.df_dataset['weights'].tolist()) - - def init_counts(self): - self.cnt_samples = 0 - self.idx_sampling = torch.multinomial(self.weights, self.num_samples, replacement=True).tolist() - - def __len__(self): - """Total number of samples of data.""" - return self.num_samples - - def __getitem__(self, idx): - if self.sampling_mode == 'whole': - new_idx = self.idx_sampling[idx] - self.cnt_samples += 1 - if self.cnt_samples >= self.num_samples: - self.init_counts() - - return self.dataset.__getitem__(new_idx) - - elif self.sampling_mode == 'half': - if torch.rand(1) > .5: - return self.dataset.__getitem__(idx) - else: - new_idx = self.idx_sampling[idx] - self.cnt_samples += 1 - if self.cnt_samples >= self.num_samples: - self.init_counts() - - return self.dataset.__getitem__(new_idx) - - elif self.sampling_mode == 'strict_half': - samples = {} - orig_sample = self.dataset.__getitem__(idx) - - new_idx = self.idx_sampling[idx] - self.cnt_samples += 1 - if self.cnt_samples >= self.num_samples: - self.init_counts() - - new_sample = self.dataset.__getitem__(new_idx) - - samples.update(orig_sample) - for k, v in new_sample.items(): - samples['os_' + k] = v - - return samples diff --git a/mpa/modules/experimental/datasets/pipelines/__init__.py b/mpa/modules/experimental/datasets/pipelines/__init__.py deleted file mode 100644 index 9dc5a0ad..00000000 --- a/mpa/modules/experimental/datasets/pipelines/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# flake8: noqa -from . import transforms diff --git a/mpa/modules/experimental/datasets/pipelines/transforms/__init__.py b/mpa/modules/experimental/datasets/pipelines/transforms/__init__.py deleted file mode 100644 index 15b668b7..00000000 --- a/mpa/modules/experimental/datasets/pipelines/transforms/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# flake8: noqa -from . import random_args_segm \ No newline at end of file diff --git a/mpa/modules/experimental/datasets/pipelines/transforms/random_args_segm.py b/mpa/modules/experimental/datasets/pipelines/transforms/random_args_segm.py deleted file mode 100644 index 440df303..00000000 --- a/mpa/modules/experimental/datasets/pipelines/transforms/random_args_segm.py +++ /dev/null @@ -1,174 +0,0 @@ -# flake8: noqa -# code in this file is adpated from -# https://github.com/ildoonet/pytorch-randaugment/blob/master/RandAugment/augmentations.py -# https://github.com/google-research/fixmatch/blob/master/third_party/auto_augment/augmentations.py -# https://github.com/google-research/fixmatch/blob/master/libml/ctaugment.py -import random - -import numpy as np -import PIL -import PIL.ImageOps -import PIL.ImageEnhance -import PIL.ImageDraw -from PIL import Image - -from mmseg.datasets.builder import PIPELINES - -PARAMETER_MAX = 10 - - -def Identity(img, **kwarg): - return img - - -def AutoContrast(img, gt, **kwarg): - return PIL.ImageOps.autocontrast(img), Identity(gt), None - - -def Brightness(img, gt, v, max_v, bias=0): - v = _float_parameter(v, max_v) + bias - return PIL.ImageEnhance.Brightness(img).enhance(v), Identity(gt), v - - -def Color(img, gt, v, max_v, bias=0): - v = _float_parameter(v, max_v) + bias - return PIL.ImageEnhance.Color(img).enhance(v), Identity(gt), v - - -def Contrast(img, gt, v, max_v, bias=0): - v = _float_parameter(v, max_v) + bias - return PIL.ImageEnhance.Contrast(img).enhance(v), Identity(gt), v - - -def Equalize(img, gt, **kwarg): - return PIL.ImageOps.equalize(img), Identity(gt), None - - -def Invert(img, gt, **kwarg): - return PIL.ImageOps.invert(img), Identity(gt), None - - -def Posterize(img, gt, v, max_v, bias=0): - v = _int_parameter(v, max_v) + bias - return PIL.ImageOps.posterize(img, v), Identity(gt), v - - -def Rotate(img, gt, v, max_v, bias=0): - v = _int_parameter(v, max_v) + bias - if random.random() < 0.5: - v = -v - return img.rotate(v), gt.rotate(v), v - - -def Sharpness(img, gt, v, max_v, bias=0): - v = _float_parameter(v, max_v) + bias - return PIL.ImageEnhance.Sharpness(img).enhance(v), Identity(gt), v - - -def ShearX(img, gt, v, max_v, bias=0): - v = _float_parameter(v, max_v) + bias - if random.random() < 0.5: - v = -v - return img.transform(img.size, PIL.Image.AFFINE, (1, v, 0, 0, 1, 0)), gt.transform(img.size, PIL.Image.AFFINE, (1, v, 0, 0, 1, 0)), v - - -def ShearY(img, gt, v, max_v, bias=0): - # print(gt) - v = _float_parameter(v, max_v) + bias - if random.random() < 0.5: - v = -v - return img.transform(img.size, PIL.Image.AFFINE, (1, 0, 0, v, 1, 0)), gt.transform(img.size, PIL.Image.AFFINE, (1, 0, 0, v, 1, 0)), v - - -def Solarize(img, gt, v, max_v, bias=0): - v = _int_parameter(v, max_v) + bias - return PIL.ImageOps.solarize(img, 256 - v), Identity(gt), v - - -def SolarizeAdd(img, gt, v, max_v, bias=0, threshold=128): - v = _int_parameter(v, max_v) + bias - if random.random() < 0.5: - v = -v - img_np = np.array(img).astype(np.int) - img_np = img_np + v - img_np = np.clip(img_np, 0, 255) - img_np = img_np.astype(np.uint8) - img = Image.fromarray(img_np) - return PIL.ImageOps.solarize(img, threshold), Identity(gt), v - - -def TranslateX(img, gt, v, max_v, bias=0): - v = _float_parameter(v, max_v) + bias - if random.random() < 0.5: - v = -v - v = int(v * img.size[0]) - return img.transform(img.size, PIL.Image.AFFINE, (1, 0, v, 0, 1, 0)), gt.transform(gt.size, PIL.Image.AFFINE, (1, 0, v, 0, 1, 0)), v - - -def TranslateY(img, gt, v, max_v, bias=0): - v = _float_parameter(v, max_v) + bias - if random.random() < 0.5: - v = -v - v = int(v * img.size[1]) - return img.transform(img.size, PIL.Image.AFFINE, (1, 0, 0, 0, 1, v)), gt.transform(gt.size, PIL.Image.AFFINE, (1, 0, 0, 0, 1, v)), v - - -def _float_parameter(v, max_v): - return float(v) * max_v / PARAMETER_MAX - - -def _int_parameter(v, max_v): - return int(v * max_v / PARAMETER_MAX) - - -def fixmatch_augment_pool(): - # FixMatch paper - augs = [(AutoContrast, None, None), - (Brightness, 0.9, 0.05), - (Color, 0.9, 0.05), - (Contrast, 0.9, 0.05), - (Equalize, None, None), - (Posterize, 4, 4), - (Rotate, 30, 0), - (Sharpness, 0.9, 0.05), - (ShearX, 0.3, 0), - (ShearY, 0.3, 0), - (Solarize, 256, 0), - (TranslateX, 0.3, 0), - (TranslateY, 0.3, 0) - ] - return augs - - -@PIPELINES.register_module() -class RandAugmentSemiSeg(object): - def __init__(self, n, m): - assert n >= 1 - assert 1 <= m <= 10 - self.n = n - self.m = m - self.augment_pool = fixmatch_augment_pool() - - def __call__(self, results): - # import time - # now = time.time() - img = results['img'] - gt = results['gt_semantic_seg'] - if not Image.isImageType(img): - img = Image.fromarray(results['img']) - if not Image.isImageType(gt): - gt = Image.fromarray(results['gt_semantic_seg'], 'L') - # img.save(f'/media/hdd2/jeom/mpa/tmp/img_{now}.png') - # gt.save(f'/media/hdd2/jeom/mpa/tmp/gt_{now}.png') - ops = random.choices(self.augment_pool, k=self.n) - for op, max_v, bias in ops: - v = np.random.randint(1, self.m) - if random.random() < 0.5: - img, gt, v = op(img, gt, v=v, max_v=max_v, bias=bias) - results['rand_semiseg_img_{}'.format(op.__name__)] = v - results['rand_semiseg_gt_{}'.format(op.__name__)] = v - # img.save(f'/media/hdd2/jeom/mpa/tmp/img_randaug_{now}.png') - # gt.save(f'/media/hdd2/jeom/mpa/tmp/gt_randaug_{now}.png') - results['img'] = np.array(img) - results['gt_semantic_seg'] = np.array(gt) - return results diff --git a/mpa/modules/experimental/datasets/pseudo_incr_dataset.py b/mpa/modules/experimental/datasets/pseudo_incr_dataset.py deleted file mode 100644 index 381d3408..00000000 --- a/mpa/modules/experimental/datasets/pseudo_incr_dataset.py +++ /dev/null @@ -1,254 +0,0 @@ -from mmdet.datasets import PIPELINES, DATASETS, CocoDataset -from mmdet.datasets.pipelines import to_tensor, MinIoURandomCrop -from mmdet.core.evaluation.bbox_overlaps import bbox_overlaps -from mmcv.parallel import DataContainer as DC - -# import torch -import numpy as np -from numpy import random - - -@DATASETS.register_module() -class PseudoIncrCocoDataset(CocoDataset): - """COCO dataset w/ pseudo label augmentation - """ - - def __init__(self, pre_stage_res, pseudo_threshold=0.9, **kwargs): - - # Build org dataset - dataset_cfg = kwargs.copy() - _ = dataset_cfg.pop('org_type', None) - super().__init__(**dataset_cfg) - - # Load pseudo labels - self.pseudo_file = pre_stage_res - self.pseudo_threshold = pseudo_threshold - self.pseudo_data = np.load(self.pseudo_file, allow_pickle=True).item() - self.pseudo_anns = self.pseudo_data['detections'] - self.pseudo_classes = self.pseudo_data['classes'] - self.CLASSES = list(self.pseudo_classes) + list(self.CLASSES) - print(f'PseudoLabelIncrDataset!!: {self.CLASSES}') - self.statistics() - - def get_ann_info(self, idx): - """Overriding CocoDataset.get_ann_info() - """ - - ann_info = super().get_ann_info(idx) - - # Adjust class labels - ann_info['labels'] += len(self.pseudo_classes) # Addin # old classes - # TODO: This could be class name sensitive adaptation - - # Format pseudo labels - pseudo_ann = self.pseudo_anns[idx] # [img][label][bbox] - labels = [] - bboxes = [] - probs = [] - for label, detections in enumerate(pseudo_ann): - for detection in detections: - if detection[4] > self.pseudo_threshold: - labels.append(label) - bboxes.append(detection[:4]) - # probabilities for old classes including BG at the end - probs.append(detection[5:]) - - # print('new_labels', ann_info['labels']) - # print('old_labels', np.array(labels)) - # print('new_bboxes', ann_info['bboxes']) - # print('old_bboxes', np.array(bboxes)) - # print(f"old: {len(labels)}, new: {ann_info['labels'].shape[0]}") - num_org_bboxes = ann_info['labels'].shape[0] - new_probs = np.zeros((num_org_bboxes, len(self.pseudo_classes) + 1), dtype=np.float32) - new_probs[:, -1] = 1.0 # NEW classes are BG for OLD classes - if len(labels) > 0: - ann_info['labels'] = \ - np.concatenate((ann_info['labels'], np.array(labels))) - ann_info['bboxes'] = \ - np.concatenate((ann_info['bboxes'], np.array(bboxes))) - ann_info['pseudo_labels'] = \ - np.concatenate((new_probs, np.array(probs))) - else: - ann_info['pseudo_labels'] = new_probs - - return ann_info - - def statistics(self): - num_bboxes = 0 - num_labels = np.zeros(len(self.pseudo_classes)) - prob_acc = np.zeros(len(self.pseudo_classes) + 1) - for pseudo_ann in self.pseudo_anns: - for label, detections in enumerate(pseudo_ann): - for detection in detections: - if detection[4] > self.pseudo_threshold: - num_bboxes += 1 - num_labels[label] += 1 - prob_acc += detection[5:] - print('pseudo label stat') - print(f'- # images: {len(self.pseudo_anns)}') - print(f'- # bboxes: {num_bboxes}') - print(f'- # labels: {num_labels}') - if num_bboxes > 0: - print(f'- label ratio: {num_labels / num_bboxes}') - print(f'- prob ratio: {prob_acc / float(num_bboxes)}') - - -@PIPELINES.register_module() -class FormatPseudoLabels(object): - """Data processor for pseudo label formatting - """ - - def __init__(self): - print('Init FormatPseudoLabels') - - def __call__(self, data): - plabels = data['pseudo_labels'] - # tensor -> DataContainer - data['pseudo_labels'] = DC(to_tensor(plabels)) - return data - - -@PIPELINES.register_module() -class LoadPseudoLabels(object): - """Data processor for pseudo label formatting - """ - - def __init__(self): - print('Init FormatPseudoLabels') - - def __call__(self, data): - plabels = data['ann_info']['pseudo_labels'] - # tensor -> DataContainer - data['pseudo_labels'] = plabels - return data - - -@PIPELINES.register_module() -class PseudoMinIoURandomCrop(MinIoURandomCrop): - """Random crop the image & bboxes, the cropped patches have minimum IoU - requirement with original image & bboxes, the IoU threshold is randomly - selected from min_ious. - - Args: - min_ious (tuple): minimum IoU threshold for all intersections with - bounding boxes - min_crop_size (float): minimum crop's size (i.e. h,w := a*h, a*w, - where a >= min_crop_size). - bbox_clip_border (bool, optional): Whether clip the objects outside - the border of the image. Defaults to True. - - Note: - The keys for bboxes, labels and masks should be paired. That is, \ - `gt_bboxes` corresponds to `gt_labels` and `gt_masks`, and \ - `gt_bboxes_ignore` to `gt_labels_ignore` and `gt_masks_ignore`. - """ - - def __init__(self, - min_ious=(0.1, 0.3, 0.5, 0.7, 0.9), - min_crop_size=0.3, - bbox_clip_border=True): - super(PseudoMinIoURandomCrop, self).__init__(min_ious, min_crop_size, bbox_clip_border) - self.bbox2label = { - 'gt_bboxes': ('gt_labels', 'pseudo_labels'), - 'gt_bboxes_ignore': 'gt_labels_ignore' - } - self.bbox2mask = { - 'gt_bboxes': ('gt_labels', 'pseudo_labels'), - 'gt_bboxes_ignore': 'gt_masks_ignore' - } - - def __call__(self, results): - """Call function to crop images and bounding boxes with minimum IoU - constraint. - - Args: - results (dict): Result dict from loading pipeline. - - Returns: - dict: Result dict with images and bounding boxes cropped, \ - 'img_shape' key is updated. - """ - - img = results['img'] - if 'bbox_fields' not in results: - raise KeyError('bbox_fields is not in results') - boxes = [results[key] for key in results['bbox_fields']] - boxes = np.concatenate(boxes, 0) - h, w, c = img.shape - while True: - mode = random.choice(self.sample_mode) - self.mode = mode - if mode == 1: - return results - - min_iou = mode - for i in range(50): - new_w = random.uniform(self.min_crop_size * w, w) - new_h = random.uniform(self.min_crop_size * h, h) - - # h / w in [0.5, 2] - if new_h / new_w < 0.5 or new_h / new_w > 2: - continue - - left = random.uniform(w - new_w) - top = random.uniform(h - new_h) - - patch = np.array( - (int(left), int(top), int(left + new_w), int(top + new_h))) - # Line or point crop is not allowed - if patch[2] == patch[0] or patch[3] == patch[1]: - continue - overlaps = bbox_overlaps( - patch.reshape(-1, 4), boxes.reshape(-1, 4)).reshape(-1) - if len(overlaps) > 0 and overlaps.min() < min_iou: - continue - - # center of boxes should inside the crop img - # only adjust boxes and instance masks when the gt is not empty - if len(overlaps) > 0: - # adjust boxes - def is_center_of_bboxes_in_patch(boxes, patch): - center = (boxes[:, :2] + boxes[:, 2:]) / 2 - mask = ((center[:, 0] > patch[0]) * - (center[:, 1] > patch[1]) * - (center[:, 0] < patch[2]) * - (center[:, 1] < patch[3])) - return mask - - mask = is_center_of_bboxes_in_patch(boxes, patch) - if not mask.any(): - continue - for key in results.get('bbox_fields', []): - boxes = results[key].copy() - mask = is_center_of_bboxes_in_patch(boxes, patch) - boxes = boxes[mask] - if self.bbox_clip_border: - boxes[:, 2:] = boxes[:, 2:].clip(max=patch[2:]) - boxes[:, :2] = boxes[:, :2].clip(min=patch[:2]) - boxes -= np.tile(patch[:2], 2) - - results[key] = boxes - # labels - label_keys = self.bbox2label.get(key) - if isinstance(label_keys, tuple): - for label_key in label_keys: - if label_key in results: - results[label_key] = results[label_key][mask] - else: - if label_keys in results: - results[label_keys] = results[label_keys][mask] - - # mask fields - mask_key = self.bbox2mask.get(key) - if mask_key in results: - results[mask_key] = results[mask_key][ - mask.nonzero()[0]].crop(patch) - # adjust the img no matter whether the gt is empty before crop - img = img[patch[1]:patch[3], patch[0]:patch[2]] - results['img'] = img - results['img_shape'] = img.shape - - # seg fields - for key in results.get('seg_fields', []): - results[key] = results[key][patch[1]:patch[3], patch[0]:patch[2]] - return results diff --git a/mpa/modules/experimental/datasets/pseudo_semi_dataset.py b/mpa/modules/experimental/datasets/pseudo_semi_dataset.py deleted file mode 100644 index 0ae98591..00000000 --- a/mpa/modules/experimental/datasets/pseudo_semi_dataset.py +++ /dev/null @@ -1,63 +0,0 @@ -from mmseg.datasets import DATASETS, build_dataset -import numpy as np - - -@DATASETS.register_module() -class PseudoSemanticSegDataset(object): - """Dataset wrapper for Semi-SL Semantic Seg experiments. - Input : splits of labeled & unlabeld datasets - """ - - def __init__(self, orig_type=None, **kwargs): - # Original dataset - dataset_cfg = kwargs.copy() - unlabeled_split = dataset_cfg.pop('unlabeled_split', '') - if unlabeled_split == '': - raise ValueError('You should use unlabeled data and put the path @ unlabeled_ann_dir!') - if 'cutmix' in dataset_cfg: - self.cutmix_flag = dataset_cfg.pop('cutmix', False) - else: - self.cutmix_flag = False - - dataset_cfg['type'] = orig_type - self.labeled_dataset = build_dataset(dataset_cfg) - self.CLASSES = self.labeled_dataset.CLASSES - self.PALETTE = self.labeled_dataset.PALETTE - - dataset_cfg['split'] = unlabeled_split - self.unlabeled_dataset = build_dataset(dataset_cfg) - - # Subsets - self.num_labeled = len(self.labeled_dataset) - self.num_unlabeled = len(self.unlabeled_dataset) - self.labeled_index = np.random.permutation(max(self.num_labeled, self.num_unlabeled)) - self.unlabeled_index = np.random.permutation(max(self.num_labeled, self.num_unlabeled)) - if self.cutmix_flag: - self.unlabeled_index2 = np.random.permutation(max(self.num_labeled, self.num_unlabeled)) - print('----------- #Labeled: ', self.num_labeled) - print('----------- #Unlabeled: ', self.num_unlabeled) - - def __len__(self): - """Total number of samples of data.""" - return self.num_labeled - - def __getitem__(self, idx): - data = {} - labeled_idx = self.labeled_index[idx] % self.num_labeled - labeled_data = self.labeled_dataset[labeled_idx] - data.update(labeled_data) - if self.num_unlabeled > 0: - tmp = self.num_unlabeled / self.num_labeled - if tmp > 0: - idx = (idx + np.random.randint(tmp)*self.num_labeled) % self.num_unlabeled - unlabeled_idx = self.unlabeled_index[idx] % self.num_unlabeled - unlabeled_data = self.unlabeled_dataset[unlabeled_idx] - for k, v in unlabeled_data.items(): - data['ul_' + k] = v - if self.cutmix_flag: - if self.num_unlabeled > 0: - unlabeled_idx = self.unlabeled_index2[idx] % self.num_unlabeled - unlabeled_data = self.unlabeled_dataset[unlabeled_idx] - for k, v in unlabeled_data.items(): - data['ul2_' + k] = v - return data diff --git a/mpa/modules/experimental/hooks/__init__.py b/mpa/modules/experimental/hooks/__init__.py deleted file mode 100644 index 7ca50e73..00000000 --- a/mpa/modules/experimental/hooks/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -# flake8: noqa -from . import stop_hook -from . import confidence_bank_hook \ No newline at end of file diff --git a/mpa/modules/experimental/hooks/confidence_bank_hook.py b/mpa/modules/experimental/hooks/confidence_bank_hook.py deleted file mode 100644 index 30c1c691..00000000 --- a/mpa/modules/experimental/hooks/confidence_bank_hook.py +++ /dev/null @@ -1,46 +0,0 @@ -import torch -import torch.nn.functional as F - -from mmcv.runner import HOOKS, Hook -# from mmcv.parallel import is_module_wrapper - - -@HOOKS.register_module() -class ConfidenceBankHook(Hook): - def __init__(self, - class_momentum=0.999, - num_classes=21, - **kwargs): - super().__init__(**kwargs) - - self.class_momentum = class_momentum - self.num_classes = num_classes - - print('####################### ConfidenceBankHook') - - def after_train_iter(self, runner): - logits = runner.outputs['logits'].detach() - logits = F.softmax(logits, dim=1) - gt = runner.outputs['gt_semantic_seg'].detach() - prev_conf = runner.data_loader._dataloader.dataset.class_criterion - - category_entropy = self.cal_category_confidence(logits, gt) - curr_conf = prev_conf * self.class_momentum + category_entropy * (1 - self.class_momentum) - - # runner.data_loader._dataloader.dataset.update_confidence_bank(curr_conf) - runner.data_loader.iter_loader._dataset.update_confidence_bank(curr_conf) - - def cal_category_confidence(self, logits, gt): - category_confidence = torch.zeros(self.num_classes).type(torch.float32) - logits = F.softmax(logits, dim=1) - for idx in range(self.num_classes): - mask_cls = (gt == idx) - if torch.sum(mask_cls) == 0: - value = 0 - else: - conf_map_sup = logits[:, idx, :, :] - value = torch.sum(conf_map_sup * mask_cls) / (torch.sum(mask_cls) + 1e-12) - - category_confidence[idx] = value - - return category_confidence diff --git a/mpa/modules/experimental/hooks/stop_hook.py b/mpa/modules/experimental/hooks/stop_hook.py deleted file mode 100644 index a1fd2202..00000000 --- a/mpa/modules/experimental/hooks/stop_hook.py +++ /dev/null @@ -1,12 +0,0 @@ -from mmcv.runner.hooks.hook import HOOKS, Hook - - -@HOOKS.register_module() -class StopHook(Hook): - def __init__(self, stop_point=15000, **kwargs): - super(StopHook, self).__init__(**kwargs) - self.stop_point = stop_point - - def before_train_iter(self, runner): - if runner.iter == self.stop_point: - runner._iter = runner._max_iters diff --git a/mpa/modules/experimental/models/__init__.py b/mpa/modules/experimental/models/__init__.py deleted file mode 100644 index 51972340..00000000 --- a/mpa/modules/experimental/models/__init__.py +++ /dev/null @@ -1,10 +0,0 @@ -# flake8: noqa -# from . import backbones -# from . import necks -from . import heads -from . import losses - -# from . import classifiers -from . import detectors -from . import segmentors -# from . import selfsl_model diff --git a/mpa/modules/experimental/models/detectors/__init__.py b/mpa/modules/experimental/models/detectors/__init__.py deleted file mode 100644 index 2a90e654..00000000 --- a/mpa/modules/experimental/models/detectors/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -# flake8: noqa -from . import lwf_two_stage_detector -from . import custom_retinanet_detector diff --git a/mpa/modules/experimental/models/detectors/custom_retinanet_detector.py b/mpa/modules/experimental/models/detectors/custom_retinanet_detector.py deleted file mode 100644 index 0c15f4ab..00000000 --- a/mpa/modules/experimental/models/detectors/custom_retinanet_detector.py +++ /dev/null @@ -1,11 +0,0 @@ -from mmdet.models.builder import DETECTORS -from mmdet.models.detectors.retinanet import RetinaNet -from mpa.modules.models.detectors.sam_detector_mixin import SAMDetectorMixin -from mpa.modules.models.detectors.l2sp_detector_mixin import L2SPDetectorMixin - - -@DETECTORS.register_module() -class CustomRetinaNet(SAMDetectorMixin, L2SPDetectorMixin, RetinaNet): - """SAM optimizer & L2SP regularizer enabled custom RetinaNet - """ - pass diff --git a/mpa/modules/experimental/models/detectors/lwf_two_stage_detector.py b/mpa/modules/experimental/models/detectors/lwf_two_stage_detector.py deleted file mode 100644 index daa4ce33..00000000 --- a/mpa/modules/experimental/models/detectors/lwf_two_stage_detector.py +++ /dev/null @@ -1,87 +0,0 @@ -from mmdet.models import DETECTORS -from mpa.modules.models.detectors.custom_two_stage_detector import CustomTwoStageDetector -import torch - - -@DETECTORS.register_module() -class LwfTwoStageDetector(CustomTwoStageDetector): - """LwF-enabled 2-stage detector - """ - - def __init__(self, arch, src_classes, dst_classes, **kwargs): - print('LwF-2-stage-detector initilaized!') - super().__init__(**kwargs) - self.arch = arch - self.src_classes = src_classes - self.dst_classes = dst_classes - - # Initialize teacher model (for OLD classes) - roi_head_cfg = kwargs.pop('roi_head', None) - if roi_head_cfg: - roi_head_cfg['bbox_head']['num_classes'] = len(src_classes) - teacher = CustomTwoStageDetector(roi_head=roi_head_cfg, **kwargs) - # - Weight fixed -> 'evaluate' mode + no gradient - teacher.eval() - # - Keep in array to prevent checkpointing teacher model - self.teachers = [teacher] - self._teacher = teacher # HOTFIX for model-level cuda()/train()/eval() - - def forward_train(self, - img, - img_metas, - gt_bboxes, - gt_labels, - gt_bboxes_ignore=None, - gt_masks=None, - proposals=None, - **kwargs): - """Learning w/o Forgetting - """ - teacher = self.teachers[0] - teacher.eval() - - # Feature - with torch.no_grad(): - x_t = teacher.extract_feat(img) - x = self.extract_feat(img) - - losses = dict() - - # Region proposal & losses - if self.with_rpn: - proposal_cfg = self.train_cfg.get('rpn_proposal', - self.test_cfg.rpn) - # Teacher (OLD tasks) - with torch.no_grad(): - rpn_outputs_t = teacher.rpn_head(x_t) - - # Student - rpn_losses, proposal_list = self.rpn_head.forward_train( - x, - img_metas, - gt_bboxes, - teacher_rpn_outputs=rpn_outputs_t, - gt_labels=None, - gt_bboxes_ignore=gt_bboxes_ignore, - proposal_cfg=proposal_cfg) - losses.update(rpn_losses) - else: - proposal_list = proposals - - # ROI outputs & losses - roi_losses = self.roi_head.forward_train( - x, - img_metas, - proposal_list, - gt_bboxes, - gt_labels, - gt_bboxes_ignore, - gt_masks, - x_t=x_t, # teacher feature map for LwF inside - roi_head_t=teacher.roi_head, # teacher roi_head for LwF inside - **kwargs) - losses.update(roi_losses) - - if self.l2sp: - losses.update(dict(loss_l2sp=self.l2sp())) - return losses diff --git a/mpa/modules/experimental/models/heads/__init__.py b/mpa/modules/experimental/models/heads/__init__.py deleted file mode 100644 index a2e84d49..00000000 --- a/mpa/modules/experimental/models/heads/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# flake8: noqa -from . import lwf_roi_head -from . import lwf_rpn_head -from . import pseudo_bbox_head -from . import lwf_bbox_head -from . import reco_sep_aspp_head diff --git a/mpa/modules/experimental/models/heads/lwf_bbox_head.py b/mpa/modules/experimental/models/heads/lwf_bbox_head.py deleted file mode 100644 index 094d2fa9..00000000 --- a/mpa/modules/experimental/models/heads/lwf_bbox_head.py +++ /dev/null @@ -1,449 +0,0 @@ -import torch -import torch.nn.functional as F - -from mmcv.runner import force_fp32 -from mmdet.models.builder import HEADS, build_loss -from mmdet.models.losses import smooth_l1_loss -from mmdet.models.dense_heads import SSDHead -from mmdet.core import multi_apply, images_to_levels, unmap, anchor_inside_flags -from mmdet.integration.nncf import no_nncf_trace - -from mpa.modules.utils.task_adapt import map_class_names - - -@HEADS.register_module() -class OffLwfSSDHead(SSDHead): - def __init__(self, src_classes, dst_classes, - loss_cls=dict(type='CrossEntropyLoss', use_sigmoid=True, loss_weight=1.0), - loss_cls_lwf=None, - **kwargs): - super().__init__(**kwargs) - self.src_classes = src_classes - self.dst_classes = dst_classes - self.loss_cls = build_loss(loss_cls) - if loss_cls_lwf is not None: - self.loss_cls_lwf = build_loss(loss_cls_lwf) - self.src2dst = torch.tensor(map_class_names(self.src_classes, self.dst_classes)) - print('OffLwfSSDHead init!') - - def forward_train(self, - x, - img_metas, - gt_bboxes, - gt_labels=None, - gt_bboxes_ignore=None, - proposal_cfg=None, - pseudo_labels=None): - """ - Args: - x (list[Tensor]): Features from FPN. - img_metas (list[dict]): Meta information of each image, e.g., - image size, scaling factor, etc. - gt_bboxes (Tensor): Ground truth bboxes of the image, - shape (num_gts, 4). - gt_labels (Tensor): Ground truth labels of each box, - shape (num_gts,). - gt_bboxes_ignore (Tensor): Ground truth bboxes to be - ignored, shape (num_ignored_gts, 4). - proposal_cfg (mmcv.Config): Test / postprocessing configuration, - if None, test_cfg would be used - - Returns: - tuple: - losses: (dict[str, Tensor]): A dictionary of loss components. - proposal_list (list[Tensor]): Proposals of each image. - """ - outs = self(x) - if gt_labels is None: - loss_inputs = outs + (gt_bboxes, img_metas) - else: - loss_inputs = outs + (gt_bboxes, gt_labels, img_metas, pseudo_labels) - with no_nncf_trace(): - losses = self.loss(*loss_inputs, gt_bboxes_ignore=gt_bboxes_ignore) - if proposal_cfg is None: - return losses - else: - proposal_list = self.get_bboxes(*outs, img_metas, cfg=proposal_cfg) - return losses, proposal_list - - def loss_single(self, cls_score, bbox_pred, anchor, labels, label_weights, - bbox_targets, bbox_weights, pseudo_labels, pseudo_label_weights, num_total_samples): - """Compute loss of a single image. - - Args: - cls_score (Tensor): Box scores for eachimage - Has shape (num_total_anchors, num_classes). - bbox_pred (Tensor): Box energies / deltas for each image - level with shape (num_total_anchors, 4). - anchors (Tensor): Box reference for each scale level with shape - (num_total_anchors, 4). - labels (Tensor): Labels of each anchors with shape - (num_total_anchors,). - label_weights (Tensor): Label weights of each anchor with shape - (num_total_anchors,) - bbox_targets (Tensor): BBox regression targets of each anchor wight - shape (num_total_anchors, 4). - bbox_weights (Tensor): BBox regression loss weights of each anchor - with shape (num_total_anchors, 4). - num_total_samples (int): If sampling, num total samples equal to - the number of total anchors; Otherwise, it is the number of - positive anchors. - - Returns: - dict[str, Tensor]: A dictionary of loss components. - """ - - loss_cls_all = F.cross_entropy( - cls_score, labels, reduction='none') * label_weights - # FG cat_id: [0, num_classes -1], BG cat_id: num_classes - pos_inds = ((labels >= 0) & - (labels < self.num_classes)).nonzero().reshape(-1) - neg_inds = (labels == self.num_classes).nonzero().view(-1) - - num_pos_samples = pos_inds.size(0) - num_neg_samples = self.train_cfg.neg_pos_ratio * num_pos_samples - if num_neg_samples > neg_inds.size(0): - num_neg_samples = neg_inds.size(0) - topk_loss_cls_neg, _ = loss_cls_all[neg_inds].topk(num_neg_samples) - loss_cls_pos = loss_cls_all[pos_inds].sum() - loss_cls_neg = topk_loss_cls_neg.sum() - loss_cls = (loss_cls_pos + loss_cls_neg) / num_total_samples - - if self.reg_decoded_bbox: - # When the regression loss (e.g. `IouLoss`, `GIouLoss`) - # is applied directly on the decoded bounding boxes, it - # decodes the already encoded coordinates to absolute format. - bbox_pred = self.bbox_coder.decode(anchor, bbox_pred) - - loss_bbox = smooth_l1_loss( - bbox_pred, - bbox_targets, - bbox_weights, - beta=self.train_cfg.smoothl1_beta, - avg_factor=num_total_samples) - loss_lwf = self.loss_cls_lwf(cls_score, pseudo_labels, - weight=pseudo_label_weights, - avg_factor=num_total_samples) - return loss_cls[None], loss_bbox, loss_lwf - - @force_fp32(apply_to=('cls_scores', 'bbox_preds')) - def loss(self, - cls_scores, - bbox_preds, - gt_bboxes, - gt_labels, - img_metas, - pseudo_labels, - gt_bboxes_ignore=None): - """Compute losses of the head. - - Args: - cls_scores (list[Tensor]): Box scores for each scale level - Has shape (N, num_anchors * num_classes, H, W) - bbox_preds (list[Tensor]): Box energies / deltas for each scale - level with shape (N, num_anchors * 4, H, W) - gt_bboxes (list[Tensor]): each item are the truth boxes for each - image in [tl_x, tl_y, br_x, br_y] format. - gt_labels (list[Tensor]): class indices corresponding to each box - img_metas (list[dict]): Meta information of each image, e.g., - image size, scaling factor, etc. - gt_bboxes_ignore (None | list[Tensor]): specify which bounding - boxes can be ignored when computing the loss. - - Returns: - dict[str, Tensor]: A dictionary of loss components. - """ - featmap_sizes = [featmap.size()[-2:] for featmap in cls_scores] - assert len(featmap_sizes) == self.anchor_generator.num_levels - - device = cls_scores[0].device - - anchor_list, valid_flag_list = self.get_anchors( - featmap_sizes, img_metas, device=device) - cls_reg_targets = self.get_targets( - anchor_list, - valid_flag_list, - gt_bboxes, - img_metas, - gt_bboxes_ignore_list=gt_bboxes_ignore, - gt_labels_list=gt_labels, - label_channels=1, - unmap_outputs=False, - pseudo_labels_list=pseudo_labels) - if cls_reg_targets is None: - return None - (labels_list, label_weights_list, bbox_targets_list, bbox_weights_list, - num_total_pos, num_total_neg, pseudo_labels_list, pseudo_label_weights_list) = cls_reg_targets - - num_images = len(img_metas) - all_cls_scores = torch.cat([ - s.permute(0, 2, 3, 1).reshape( - num_images, -1, self.cls_out_channels) for s in cls_scores - ], 1) - all_labels = torch.cat(labels_list, -1).view(num_images, -1) - all_label_weights = torch.cat(label_weights_list, - -1).view(num_images, -1) - all_bbox_preds = torch.cat([ - b.permute(0, 2, 3, 1).reshape(num_images, -1, 4) - for b in bbox_preds - ], -2) - all_bbox_targets = torch.cat(bbox_targets_list, - -2).view(num_images, -1, 4) - all_bbox_weights = torch.cat(bbox_weights_list, - -2).view(num_images, -1, 4) - all_pseudo_labels = torch.cat(pseudo_labels_list, - -2).view(num_images, -1, len(self.src_classes) + 1) - all_pseudo_label_weights = torch.cat(pseudo_label_weights_list, - -1).view(num_images, -1, 1) - - # concat all level anchors to a single tensor - all_anchors = [] - for i in range(num_images): - all_anchors.append(torch.cat(anchor_list[i])) - - # check NaN and Inf - assert torch.isfinite(all_cls_scores).all().item(), \ - 'classification scores become infinite or NaN!' - assert torch.isfinite(all_bbox_preds).all().item(), \ - 'bbox predications become infinite or NaN!' - - losses_cls, losses_bbox, losses_lwf = multi_apply( - self.loss_single, - all_cls_scores, - all_bbox_preds, - all_anchors, - all_labels, - all_label_weights, - all_bbox_targets, - all_bbox_weights, - all_pseudo_labels, - all_pseudo_label_weights, - num_total_samples=num_total_pos) - - if self.loss_balancing: - losses_cls, losses_reg, losses_lwf = self._balance_losses(losses_cls, - losses_bbox, - losses_lwf) - - return dict(loss_cls=losses_cls, loss_bbox=losses_bbox, loss_lwf=losses_lwf) - - def _balance_losses(self, losses_cls, losses_reg, losses_lwf): - loss_cls = sum(_loss.mean() for _loss in losses_cls) - loss_cls = torch.exp(-self.loss_weights[0]) * loss_cls + 0.5 * self.loss_weights[0] - - loss_reg = sum(_loss.mean() for _loss in losses_reg) - loss_reg = torch.exp(-self.loss_weights[1]) * loss_reg + 0.5 * self.loss_weights[1] - - loss_lwf = sum(_loss.mean() for _loss in losses_lwf) - loss_lwf = torch.exp(-self.loss_weights[1]) * loss_lwf + 0.5 * self.loss_weights[1] - - return loss_cls, loss_reg, loss_lwf - - def _get_targets_single(self, - flat_anchors, - valid_flags, - gt_bboxes, - gt_bboxes_ignore, - gt_labels, - pseudo_labels, - img_meta, - label_channels=1, - unmap_outputs=True): - """Compute regression and classification targets for anchors in a - single image. - - Args: - flat_anchors (Tensor): Multi-level anchors of the image, which are - concatenated into a single tensor of shape (num_anchors ,4) - valid_flags (Tensor): Multi level valid flags of the image, - which are concatenated into a single tensor of - shape (num_anchors,). - gt_bboxes (Tensor): Ground truth bboxes of the image, - shape (num_gts, 4). - img_meta (dict): Meta info of the image. - gt_bboxes_ignore (Tensor): Ground truth bboxes to be - ignored, shape (num_ignored_gts, 4). - img_meta (dict): Meta info of the image. - gt_labels (Tensor): Ground truth labels of each box, - shape (num_gts,). - label_channels (int): Channel of label. - unmap_outputs (bool): Whether to map outputs back to the original - set of anchors. - - Returns: - tuple: - labels_list (list[Tensor]): Labels of each level - label_weights_list (list[Tensor]): Label weights of each level - bbox_targets_list (list[Tensor]): BBox targets of each level - bbox_weights_list (list[Tensor]): BBox weights of each level - num_total_pos (int): Number of positive samples in all images - num_total_neg (int): Number of negative samples in all images - """ - inside_flags = anchor_inside_flags(flat_anchors, valid_flags, - img_meta['img_shape'][:2], - self.train_cfg.allowed_border) - if not inside_flags.any(): - return (None,) * 7 - # assign gt and sample anchors - anchors = flat_anchors[inside_flags, :] - - assign_result = self.assigner.assign( - anchors, gt_bboxes, gt_bboxes_ignore, - None if self.sampling else gt_labels) - sampling_result = self.sampler.sample(assign_result, anchors, - gt_bboxes) - - num_valid_anchors = anchors.shape[0] - bbox_targets = torch.zeros_like(anchors) - bbox_weights = torch.zeros_like(anchors) - labels = anchors.new_full((num_valid_anchors,), - self.num_classes, - dtype=torch.long) - label_weights = anchors.new_zeros(num_valid_anchors, dtype=torch.float) - plabels = anchors.new_zeros((num_valid_anchors, len(self.src_classes) + 1), dtype=torch.float) - plabel_weights = anchors.new_zeros(num_valid_anchors, dtype=torch.float) - - pos_inds = sampling_result.pos_inds - neg_inds = sampling_result.neg_inds - if len(pos_inds) > 0: - if not self.reg_decoded_bbox: - pos_bbox_targets = self.bbox_coder.encode( - sampling_result.pos_bboxes, sampling_result.pos_gt_bboxes) - else: - pos_bbox_targets = sampling_result.pos_gt_bboxes - bbox_targets[pos_inds, :] = pos_bbox_targets - bbox_weights[pos_inds, :] = 1.0 - if gt_labels is None: - # Only rpn gives gt_labels as None - # Foreground is the first class since v2.5.0 - labels[pos_inds] = 0 - else: - labels[pos_inds] = gt_labels[ - sampling_result.pos_assigned_gt_inds] - if self.train_cfg.pos_weight <= 0: - label_weights[pos_inds] = 1.0 - else: - label_weights[pos_inds] = self.train_cfg.pos_weight - if pseudo_labels is not None: - plabels[pos_inds] = pseudo_labels[sampling_result.pos_assigned_gt_inds] - plabel_weights[pos_inds] = 1.0 - if len(neg_inds) > 0: - label_weights[neg_inds] = 1.0 - - # map up to original set of anchors - if unmap_outputs: - num_total_anchors = flat_anchors.size(0) - labels = unmap( - labels, num_total_anchors, inside_flags, - fill=self.num_classes) # fill bg label - label_weights = unmap(label_weights, num_total_anchors, - inside_flags) - bbox_targets = unmap(bbox_targets, num_total_anchors, inside_flags) - bbox_weights = unmap(bbox_weights, num_total_anchors, inside_flags) - plabels = unmap(plabels, num_total_anchors, inside_flags) - plabel_weights = unmap(plabel_weights, num_total_anchors, inside_flags) - return (labels, label_weights, bbox_targets, bbox_weights, pos_inds, - neg_inds, sampling_result, plabels, plabel_weights) - - def get_targets(self, - anchor_list, - valid_flag_list, - gt_bboxes_list, - img_metas, - gt_bboxes_ignore_list=None, - gt_labels_list=None, - pseudo_labels_list=None, - label_channels=1, - unmap_outputs=True, - return_sampling_results=False): - """Compute regression and classification targets for anchors in - multiple images. - - Args: - anchor_list (list[list[Tensor]]): Multi level anchors of each - image. The outer list indicates images, and the inner list - corresponds to feature levels of the image. Each element of - the inner list is a tensor of shape (num_anchors, 4). - valid_flag_list (list[list[Tensor]]): Multi level valid flags of - each image. The outer list indicates images, and the inner list - corresponds to feature levels of the image. Each element of - the inner list is a tensor of shape (num_anchors, ) - gt_bboxes_list (list[Tensor]): Ground truth bboxes of each image. - img_metas (list[dict]): Meta info of each image. - gt_bboxes_ignore_list (list[Tensor]): Ground truth bboxes to be - ignored. - gt_labels_list (list[Tensor]): Ground truth labels of each box. - label_channels (int): Channel of label. - unmap_outputs (bool): Whether to map outputs back to the original - set of anchors. - - Returns: - tuple: Usually returns a tuple containing learning targets. - - - labels_list (list[Tensor]): Labels of each level. - - label_weights_list (list[Tensor]): Label weights of each \ - level. - - bbox_targets_list (list[Tensor]): BBox targets of each level. - - bbox_weights_list (list[Tensor]): BBox weights of each level. - - num_total_pos (int): Number of positive samples in all \ - images. - - num_total_neg (int): Number of negative samples in all \ - images. - additional_returns: This function enables user-defined returns from - `self._get_targets_single`. These returns are currently refined - to properties at each feature map (i.e. having HxW dimension). - The results will be concatenated after the end - """ - num_imgs = len(img_metas) - assert len(anchor_list) == len(valid_flag_list) == num_imgs - # anchor number of multi levels - num_level_anchors = [anchors.size(0) for anchors in anchor_list[0]] - # concat all level anchors to a single tensor - concat_anchor_list = [] - concat_valid_flag_list = [] - for i in range(num_imgs): - assert len(anchor_list[i]) == len(valid_flag_list[i]) - concat_anchor_list.append(torch.cat(anchor_list[i])) - concat_valid_flag_list.append(torch.cat(valid_flag_list[i])) - - # compute targets for each image - if gt_bboxes_ignore_list is None: - gt_bboxes_ignore_list = [None for _ in range(num_imgs)] - if gt_labels_list is None: - gt_labels_list = [None for _ in range(num_imgs)] - results = multi_apply( - self._get_targets_single, - concat_anchor_list, - concat_valid_flag_list, - gt_bboxes_list, - gt_bboxes_ignore_list, - gt_labels_list, - pseudo_labels_list, - img_metas, - label_channels=label_channels, - unmap_outputs=unmap_outputs) - (all_labels, all_label_weights, all_bbox_targets, all_bbox_weights, - pos_inds_list, neg_inds_list, sampling_results_list) = results[:7] - rest_results = list(results[7:]) # user-added return values - # no valid anchors - if any([labels is None for labels in all_labels]): - return None - # sampled anchors of all images - num_total_pos = sum([max(inds.numel(), 1) for inds in pos_inds_list]) - num_total_neg = sum([max(inds.numel(), 1) for inds in neg_inds_list]) - # split targets to a list w.r.t. multiple levels - labels_list = images_to_levels(all_labels, num_level_anchors) - label_weights_list = images_to_levels(all_label_weights, - num_level_anchors) - bbox_targets_list = images_to_levels(all_bbox_targets, - num_level_anchors) - bbox_weights_list = images_to_levels(all_bbox_weights, - num_level_anchors) - res = (labels_list, label_weights_list, bbox_targets_list, - bbox_weights_list, num_total_pos, num_total_neg) - if return_sampling_results: - res = res + (sampling_results_list,) - for i, r in enumerate(rest_results): # user-added return values - rest_results[i] = images_to_levels(r, num_level_anchors) - - return res + tuple(rest_results) diff --git a/mpa/modules/experimental/models/heads/lwf_roi_head.py b/mpa/modules/experimental/models/heads/lwf_roi_head.py deleted file mode 100644 index 5199adc7..00000000 --- a/mpa/modules/experimental/models/heads/lwf_roi_head.py +++ /dev/null @@ -1,232 +0,0 @@ -import torch - -from mmcv.runner import force_fp32 -from mmdet.models import HEADS, build_loss -from mmdet.models.roi_heads.standard_roi_head import StandardRoIHead -from mmdet.core.bbox.transforms import bbox2roi - -from mpa.modules.utils.task_adapt import map_class_names - - -@HEADS.register_module() -class LwfRoIHead(StandardRoIHead): - """LwF-enabled RoI head - """ - - def __init__(self, src_classes, dst_classes, loss_cls_lwf, loss_bbox_lwf, **kwargs): - super().__init__(**kwargs) - self.src_classes = src_classes - self.dst_classes = dst_classes - self.loss_cls_lwf = build_loss(loss_cls_lwf) - self.loss_bbox_lwf = build_loss(loss_bbox_lwf) - self.src2dst = torch.tensor(map_class_names(self.src_classes, self.dst_classes)) - print('LwfRoIHead init!') - - def forward_train(self, - x, - img_metas, - proposal_list, - gt_bboxes, - gt_labels, - gt_bboxes_ignore=None, - gt_masks=None, - x_t=None, - roi_head_t=None): - """ - Args: - x_t (list[Tensor]): feature maps from teacher model - roid_head_t (StandardRoIHead): feature RoI head for LwF - - Returns: - dict[str, Tensor]: a dictionary of loss components - """ - # assign gts and sample proposals - if self.with_bbox or self.with_mask: - num_imgs = len(img_metas) - if gt_bboxes_ignore is None: - gt_bboxes_ignore = [None for _ in range(num_imgs)] - sampling_results = [] - for i in range(num_imgs): - assign_result = self.bbox_assigner.assign( - proposal_list[i], gt_bboxes[i], gt_bboxes_ignore[i], - gt_labels[i]) - sampling_result = self.bbox_sampler.sample( - assign_result, - proposal_list[i], - gt_bboxes[i], - gt_labels[i], - feats=[lvl_feat[i][None] for lvl_feat in x]) - sampling_results.append(sampling_result) - - losses = dict() - # bbox head forward and loss - if self.with_bbox: - # Student - bbox_results = self._bbox_forward_train(x, sampling_results, - gt_bboxes, gt_labels, - img_metas) - # Normal BBox losses - losses.update(bbox_results['loss_bbox']) - - # Teacher - with torch.no_grad(): - bbox_results_t = roi_head_t._bbox_forward(x_t, bbox_results['rois']) - - # LwF losses - lwf_losses = self.loss_lwf( - bbox_results['labels'], - bbox_results['cls_score'], bbox_results['bbox_pred'], - bbox_results_t['cls_score'], bbox_results_t['bbox_pred']) - losses.update(lwf_losses) - - # mask head forward and loss - if self.with_mask: - mask_results = self._mask_forward_train(x, sampling_results, - bbox_results['bbox_feats'], - gt_masks, img_metas) - losses.update(mask_results['loss_mask']) - - return losses - - @force_fp32(apply_to=('cls_scores', 'bbox_preds', 'cls_scores_t', 'bbox_preds_t')) - def loss_lwf(self, gt_labels, cls_score, bbox_pred, cls_score_t, bbox_pred_t): - # Classification LwF loss - # Only for BG-labeled bboxes - is_bg_labels = (gt_labels == len(self.dst_classes)) - cls_score = cls_score[is_bg_labels] - cls_score_t = cls_score_t[is_bg_labels] - cls_loss = self.loss_cls_lwf(cls_score, cls_score_t) - - # Regression LwF loss - if not self.bbox_head.reg_class_agnostic: - n = bbox_pred.shape[0] - bbox_pred = bbox_pred.view(n, -1, 4)[:, self.src2dst[self.src2dst >= 0]] - bbox_pred_t = bbox_pred_t.view(n, -1, 4)[:, self.src2dst >= 0] - - bbox_loss = self.loss_bbox_lwf(bbox_pred, bbox_pred_t) - return dict(loss_cls_lwf=cls_loss, loss_bbox_lwf=bbox_loss) - - def _bbox_forward_train(self, x, sampling_results, gt_bboxes, gt_labels, - img_metas): - rois = bbox2roi([res.bboxes for res in sampling_results]) - bbox_results = self._bbox_forward(x, rois) - - bbox_targets = self.bbox_head.get_targets(sampling_results, gt_bboxes, - gt_labels, self.train_cfg) - loss_bbox = self.bbox_head.loss(bbox_results['cls_score'], - bbox_results['bbox_pred'], rois, - *bbox_targets) - - bbox_results.update( - loss_bbox=loss_bbox, - rois=rois, - labels=bbox_targets[0], - label_weights=bbox_targets[1], - bbox_targets=bbox_targets[2], - bbox_weights=bbox_targets[3] - ) - return bbox_results - - -@HEADS.register_module() -class OffLwfRoIHead(StandardRoIHead): - """Offline LwF-enabled RoI head - """ - - def __init__(self, src_classes, dst_classes, loss_cls_lwf, **kwargs): - super().__init__(**kwargs) - self.src_classes = src_classes - self.dst_classes = dst_classes - self.loss_cls_lwf = build_loss(loss_cls_lwf) - self.src2dst = torch.tensor(map_class_names(self.src_classes, self.dst_classes)) - print('OffLwfRoIHead init!') - - def forward_train(self, - x, - img_metas, - proposal_list, - gt_bboxes, - gt_labels, - gt_bboxes_ignore=None, - gt_masks=None, - pseudo_labels=None): - """ - Args: - pseudo_labels (list[Tensor]): class proabilities for OLD classes - - Returns: - dict[str, Tensor]: a dictionary of loss components - """ - # assign gts and sample proposals - if self.with_bbox or self.with_mask: - num_imgs = len(img_metas) - if gt_bboxes_ignore is None: - gt_bboxes_ignore = [None for _ in range(num_imgs)] - sampling_results = [] - for i in range(num_imgs): - assign_result = self.bbox_assigner.assign( - proposal_list[i], gt_bboxes[i], gt_bboxes_ignore[i], - gt_labels[i]) - sampling_result = self.bbox_sampler.sample( - assign_result, - proposal_list[i], - gt_bboxes[i], - gt_labels[i], - feats=[lvl_feat[i][None] for lvl_feat in x]) - sampling_results.append(sampling_result) - - losses = dict() - # bbox head forward and loss - if self.with_bbox: - bbox_results = self._bbox_forward_train(x, sampling_results, - gt_bboxes, gt_labels, - pseudo_labels, # Only one difference from super().forward_train() - img_metas) - losses.update(bbox_results['loss_bbox']) - - # mask head forward and loss - if self.with_mask: - mask_results = self._mask_forward_train(x, sampling_results, - bbox_results['bbox_feats'], - gt_masks, img_metas) - losses.update(mask_results['loss_mask']) - - return losses - - def _bbox_forward_train(self, x, sampling_results, gt_bboxes, gt_labels, pseudo_labels, - img_metas): - """Run forward function and calculate loss for box head in training.""" - rois = bbox2roi([res.bboxes for res in sampling_results]) - bbox_results = self._bbox_forward(x, rois) - - target_labels, label_weights, target_bboxes, bbox_weights = \ - self.bbox_head.get_targets(sampling_results, gt_bboxes, gt_labels, self.train_cfg) - - # Pseudo label targets - old_plabels = [] - old_pred_indices = [] - num_targets = 0 - for img_idx, sampling_result in enumerate(sampling_results): - pos_plabels = pseudo_labels[img_idx][sampling_result.pos_assigned_gt_inds] - old_indices = torch.nonzero(pos_plabels[:, -1] < 1.0, as_tuple=False)[:, -1] # OLD class pseudo labels - pos_plabels = pos_plabels[old_indices] - old_plabels.append(pos_plabels) - old_pred_indices.append(old_indices + num_targets) - num_targets += sampling_result.bboxes.shape[0] - old_plabels = torch.cat(old_plabels) - old_pred_indices = torch.cat(old_pred_indices) - - label_weights[old_pred_indices] = 0.0 # Suppress x-entropy for OLD predictions - - # Ordinary NEW class loss - loss_bbox = self.bbox_head.loss(bbox_results['cls_score'], - bbox_results['bbox_pred'], rois, - target_labels, label_weights, target_bboxes, bbox_weights) - bbox_results.update(loss_bbox=loss_bbox) - - # LwF loss for OLD classes - loss_cls_lwf = self.loss_cls_lwf( - bbox_results['cls_score'][old_pred_indices], old_plabels, target_is_logit=False) - bbox_results['loss_bbox'].update(loss_cls_lwf=loss_cls_lwf) - - return bbox_results diff --git a/mpa/modules/experimental/models/heads/lwf_rpn_head.py b/mpa/modules/experimental/models/heads/lwf_rpn_head.py deleted file mode 100644 index 70865114..00000000 --- a/mpa/modules/experimental/models/heads/lwf_rpn_head.py +++ /dev/null @@ -1,89 +0,0 @@ -import torch - -from mmcv.runner import force_fp32 -from mmdet.models import HEADS, build_loss -from mmdet.models.dense_heads import RPNHead -from mmdet.core import multi_apply - - -@HEADS.register_module() -class LwfRPNHead(RPNHead): - """LwF-enabled RPN head - """ - - def __init__(self, in_channels, loss_cls_lwf, loss_bbox_lwf, **kwargs): - super().__init__(in_channels, **kwargs) - self.loss_cls_lwf = build_loss(loss_cls_lwf) - self.loss_bbox_lwf = build_loss(loss_bbox_lwf) - print('LwfRPNHead init!') - - def forward_train(self, - x, - img_metas, - gt_bboxes, - teacher_rpn_outputs=None, - gt_labels=None, - gt_bboxes_ignore=None, - proposal_cfg=None, - **kwargs): - - # Normal losses & proposals - outputs = self(x) - if gt_labels is None: - loss_inputs = outputs + (gt_bboxes, img_metas) - else: - loss_inputs = outputs + (gt_bboxes, gt_labels, img_metas) - losses = self.loss(*loss_inputs, gt_bboxes_ignore=gt_bboxes_ignore) - - # Compute LwF losses - lwf_losses = self.loss_lwf(*outputs, *teacher_rpn_outputs) - losses.update(lwf_losses) - - # Proposals - if proposal_cfg is None: - return losses - else: - proposal_list = self.get_bboxes(*outputs, img_metas, cfg=proposal_cfg) - return losses, proposal_list - - @force_fp32(apply_to=('cls_scores', 'bbox_preds', 'cls_scores_t', 'bbox_preds_t')) - def loss_lwf(self, cls_scores, bbox_preds, cls_scores_t, bbox_preds_t): - """Compute multi-scale LwF losses - - Args: - cls_scores (list[Tensor]): Box scores for each scale level - Has shape (N, num_anchors * num_classes, H, W) - bbox_preds (list[Tensor]): Box energies / deltas for each scale - level with shape (N, num_anchors * 4, H, W) - cls_scores_t (list[Tensor]): Box scores of teacher model - bbox_preds_t (list[Tensor]): Box energies / deltas of teacher model - - Returns: - dict[str, Tensor]: A dictionary of loss components. - """ - cls_losses, bbox_losses = multi_apply( - self.loss_lwf_single, - cls_scores, - bbox_preds, - cls_scores_t, - bbox_preds_t - ) - return dict(loss_rpn_cls_lwf=cls_losses, loss_rpn_bbox_lwf=bbox_losses) - - def loss_lwf_single(self, cls_score, bbox_pred, cls_score_t, bbox_pred_t): - """Compute single-scale LwF losses - """ - # Classification LwF loss - cls_score = cls_score.permute(0, 2, 3, 1).reshape(-1, self.cls_out_channels) - cls_score_t = cls_score_t.permute(0, 2, 3, 1).reshape(-1, self.cls_out_channels) - cls_prob_t = torch.sigmoid(cls_score_t) - cls_loss = self.loss_cls_lwf(cls_score, cls_score_t, - weight=cls_prob_t) # More weights on teacher's positives - - # Regression LwF loss - bbox_pred = bbox_pred.permute(0, 2, 3, 1).reshape(-1, 4) - bbox_pred_t = bbox_pred_t.permute(0, 2, 3, 1).reshape(-1, 4) - bbox_loss = self.loss_bbox_lwf(bbox_pred, bbox_pred_t, - weight=cls_prob_t) # More weights on teacher's positives - - return cls_loss, bbox_loss diff --git a/mpa/modules/experimental/models/heads/pseudo_bbox_head.py b/mpa/modules/experimental/models/heads/pseudo_bbox_head.py deleted file mode 100644 index ca44075f..00000000 --- a/mpa/modules/experimental/models/heads/pseudo_bbox_head.py +++ /dev/null @@ -1,113 +0,0 @@ -import torch -import torch.nn.functional as F - -from mmcv.runner import force_fp32 -from mmdet.core.utils.misc import topk -from mmdet.models.builder import HEADS -from mmdet.models.roi_heads import Shared2FCBBoxHead -from mmdet.models.dense_heads import SSDHead - -from mpa.modules.experimental.utils.pseudo_nms import pseudo_multiclass_nms - - -@HEADS.register_module() -class PseudoShared2FCBBoxHead(Shared2FCBBoxHead): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - - @force_fp32(apply_to=('cls_score', 'bbox_pred')) - def get_bboxes(self, - rois, - cls_score, - bbox_pred, - img_shape, - scale_factor, - rescale=False, - cfg=None): - """Transform model outputs to bboxes after NMS - - Copied from mmdet/models/roi_heads/bbox_heads/bbox_head.py - to augment detection output with class probabilities as pseudo label - """ - if isinstance(cls_score, list): - cls_score = sum(cls_score) / float(len(cls_score)) - scores = F.softmax(cls_score, dim=1) if cls_score is not None else None - - if bbox_pred is not None: - bboxes = self.bbox_coder.decode( - rois[:, 1:], bbox_pred, max_shape=img_shape) - else: - bboxes = rois[:, 1:].clone() - if img_shape is not None: - bboxes[:, [0, 2]].clamp_(min=0, max=img_shape[1]) - bboxes[:, [1, 3]].clamp_(min=0, max=img_shape[0]) - - if rescale and bboxes.size(0) > 0: - if isinstance(scale_factor, float): - bboxes /= scale_factor - else: - scale_factor = bboxes.new_tensor(scale_factor) - bboxes = (bboxes.view(bboxes.size(0), -1, 4) / - scale_factor).view(bboxes.size()[0], -1) - - if cfg is None: - return bboxes, scores - else: - det_bboxes, det_labels = pseudo_multiclass_nms(bboxes, scores, - cfg.score_thr, cfg.nms, - cfg.max_per_img) - - return det_bboxes, det_labels - - -@HEADS.register_module() -class PseudoSSDHead(SSDHead): - - def _get_bboxes_single(self, - cls_score_list, - bbox_pred_list, - mlvl_anchors, - img_shape, - scale_factor, - cfg, - rescale=False): - """ - Transform outputs for a single batch item into labeled boxes. - - Copied from mmdet/models/dense_heads/anchor_head.py - to augment detection output with class probabilities as pseudo label - """ - cfg = self.test_cfg if cfg is None else cfg - assert len(cls_score_list) == len(bbox_pred_list) == len(mlvl_anchors) - mlvl_bboxes = [] - mlvl_scores = [] - for cls_score, bbox_pred, anchors in zip(cls_score_list, - bbox_pred_list, mlvl_anchors): - assert cls_score.size()[-2:] == bbox_pred.size()[-2:] - cls_score = cls_score.permute(1, 2, - 0).reshape(-1, self.cls_out_channels) - if self.use_sigmoid_cls: - scores = cls_score.sigmoid() - else: - scores = cls_score.softmax(-1) - # scores = scores[:, :self.num_classes] # Removed this line to retain BG prob - bbox_pred = bbox_pred.permute(1, 2, 0).reshape(-1, 4) - nms_pre = cfg.get('nms_pre', -1) - if nms_pre > 0: - max_scores, _ = scores.max(dim=1) - _, topk_inds = topk(max_scores, nms_pre) - anchors = anchors[topk_inds, :] - bbox_pred = bbox_pred[topk_inds, :] - scores = scores[topk_inds, :] - bboxes = self.bbox_coder.decode( - anchors, bbox_pred, max_shape=img_shape) - mlvl_bboxes.append(bboxes) - mlvl_scores.append(scores) - mlvl_bboxes = torch.cat(mlvl_bboxes) - if rescale: - mlvl_bboxes /= mlvl_bboxes.new_tensor(scale_factor) - mlvl_scores = torch.cat(mlvl_scores) - det_bboxes, det_labels = pseudo_multiclass_nms(mlvl_bboxes, mlvl_scores, # Replaced multiclass_nms by pseudo_* - cfg.score_thr, cfg.nms, - cfg.max_per_img) - return det_bboxes, det_labels diff --git a/mpa/modules/experimental/models/heads/reco_sep_aspp_head.py b/mpa/modules/experimental/models/heads/reco_sep_aspp_head.py deleted file mode 100644 index fd5a79af..00000000 --- a/mpa/modules/experimental/models/heads/reco_sep_aspp_head.py +++ /dev/null @@ -1,187 +0,0 @@ -import torch -import torch.nn as nn -from mmcv.cnn import DepthwiseSeparableConvModule - -from mmseg.ops import resize -from mmseg.models.builder import HEADS -from mmseg.models.decode_heads.sep_aspp_head import DepthwiseSeparableASPPHead -from mmseg.models.losses import accuracy -from mmseg.models.builder import build_loss -from mmcv.runner import force_fp32 - - -@HEADS.register_module() -class ReCoDepthwiseSeparableASPPHead(DepthwiseSeparableASPPHead): - """Encoder-Decoder with Atrous Separable Convolution for Semantic Image - Segmentation. - - This head is the implementation of `DeepLabV3+ - `_. - - Args: - c1_in_channels (int): The input channels of c1 decoder. If is 0, - the no decoder will be used. - c1_channels (int): The intermediate channels of c1 decoder. - """ - - def __init__(self, reco_loss_decode=dict(type='ReCoLoss', loss_weight=1.0), **kwargs): - super(ReCoDepthwiseSeparableASPPHead, self).__init__(**kwargs) - c1_channels = kwargs['c1_channels'] - self.reco_loss_decode = build_loss(reco_loss_decode) - # reco HP - self.weak_threshold = 0.7 - self.strong_threshold = 0.97 - self.temp = 0.5 - self.num_queries = 256 - self.num_negatives = 256 - - self.sep_representation = nn.Sequential( - DepthwiseSeparableConvModule( - self.channels + c1_channels, - self.channels, - 3, - padding=1, - norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg), - DepthwiseSeparableConvModule( - self.channels, - self.channels, - 3, - padding=1, - norm_cfg=self.norm_cfg, - act_cfg=self.act_cfg), - nn.Conv2d(self.channels, 256, kernel_size=1)) - - def forward(self, inputs): - """Forward function.""" - x = self._transform_inputs(inputs) - aspp_outs = [ - resize(self.image_pool(x), - size=x.size()[2:], - mode='bilinear', - align_corners=self.align_corners)] - aspp_outs.extend(self.aspp_modules(x)) - aspp_outs = torch.cat(aspp_outs, dim=1) - output = self.bottleneck(aspp_outs) - if self.c1_bottleneck is not None: - c1_output = self.c1_bottleneck(inputs[0]) - output = resize( - input=output, - size=c1_output.shape[2:], - mode='bilinear', - align_corners=self.align_corners) - output = torch.cat([output, c1_output], dim=1) - representation = self.sep_representation(output) - output = self.sep_bottleneck(output) - output = self.cls_seg(output) - return output, representation - - def ce_loss(self, seg_logit, gt, weight=None): - if self.sampler is not None: - seg_weight = self.sampler.sample(seg_logit, gt) - else: - if weight is None: - seg_weight = weight - else: - batch_size = seg_logit.shape[0] - valid_mask = (gt < 255).float() - weight_view = weight.view(batch_size, -1).ge(self.strong_threshold).sum(-1) - weighting = weight_view / valid_mask.view(batch_size, -1).sum(-1) - # print("weighting:{}".format(weighting)) - weighting = weighting.view(batch_size, 1, 1) - seg_weight = weighting - - seg_label = gt.squeeze(1) - loss_seg = self.loss_decode( - seg_logit, - seg_label, - weight=seg_weight, - ignore_index=self.ignore_index - ) - acc_seg = accuracy(seg_logit, seg_label) - - return loss_seg, acc_seg - - def unsupervised_reco_loss(self, rep, label, mask, prob, - strong_threshold=1.0, temp=0.5, - num_queries=256, num_negatives=256): - - seg_label = label.squeeze(1) - loss_seg = self.reco_loss_decode(rep, label, mask, prob, - strong_threshold=strong_threshold, temp=temp, - num_queries=num_queries, num_negatives=num_negatives) - acc_seg = accuracy(prob, seg_label) - - return loss_seg, acc_seg - - @force_fp32(apply_to=('seg_logit', 'seg_u_logits',)) - def reco_loss(self, feature, gt_semantic_seg, mask, logits, - strong_threshold=0.97, temp=0.5, - num_queries=256, num_negatives=256): - """Compute segmentation loss.""" - loss = dict() - gt = gt_semantic_seg['gt'] - pseudo_label = gt_semantic_seg['pseudo_label'] - seg_logits, seg_u_logits = logits - train_u_reco_mask, train_u_pseudo_mask = mask - - seg_logit = resize(input=seg_logits, - size=gt.shape[2:], - mode='bilinear', - align_corners=self.align_corners) - seg_u_logit = resize(input=seg_u_logits, - size=pseudo_label.shape[2:], - mode='bilinear', - align_corners=self.align_corners) - loss['loss_seg'], loss['acc_seg'] = self.ce_loss(seg_logit, gt) - loss['loss_seg_u'], loss['acc_seg_u'] = self.ce_loss(seg_u_logit, pseudo_label, train_u_pseudo_mask) - pseudo_label = resize(input=pseudo_label.float(), size=feature.shape[2:], mode='nearest') - train_u_reco_mask = resize(input=train_u_reco_mask.unsqueeze(1), size=feature.shape[2:], mode='nearest') - loss['loss_seg_reco'], loss['acc_reco'] = self.unsupervised_reco_loss(feature, pseudo_label, - train_u_reco_mask, - seg_u_logits, - strong_threshold=strong_threshold, - temp=temp, - num_queries=num_queries, - num_negatives=num_negatives) - # loss['loss'] = loss['loss_seg'] + loss['loss_seg_u'] + loss['loss_seg_reco'] - - return loss - - def forward_train(self, inputs, img_metas, gt_semantic_seg, train_cfg): - if isinstance(inputs, dict): - x = inputs['x'] - x_u = inputs['x_u'] - conf = inputs['conf'] - - seg_logits, _ = self.forward(x) - seg_u_logits, feature = self.forward(x_u) - logits = (seg_logits, seg_u_logits) - train_u_reco_mask = conf.ge(self.weak_threshold).float() - train_u_pseudo_mask = conf.ge(self.strong_threshold).float() - mask = (train_u_reco_mask, train_u_pseudo_mask) - losses = self.reco_loss(feature, gt_semantic_seg, mask, logits, - strong_threshold=self.strong_threshold, temp=self.temp, - num_queries=self.num_queries, num_negatives=self.num_negatives) - else: - seg_logits = self.forward(inputs) - losses = self.losses(seg_logits, gt_semantic_seg) - return losses - - def forward_test(self, inputs, img_metas, test_cfg): - """Forward function for testing. - - Args: - inputs (list[Tensor]): List of multi-level img features. - img_metas (list[dict]): List of image info dict where each dict - has: 'img_shape', 'scale_factor', 'flip', and may also contain - 'filename', 'ori_shape', 'pad_shape', and 'img_norm_cfg'. - For details on the values of these keys see - `mmseg/datasets/pipelines/formatting.py:Collect`. - test_cfg (dict): The testing config. - - Returns: - Tensor: Output segmentation map. - """ - seg_logits, _ = self.forward(inputs) - return seg_logits diff --git a/mpa/modules/experimental/models/losses/__init__.py b/mpa/modules/experimental/models/losses/__init__.py deleted file mode 100644 index 47e16a40..00000000 --- a/mpa/modules/experimental/models/losses/__init__.py +++ /dev/null @@ -1,11 +0,0 @@ -# flake8: noqa -from . import distillation_losses -from . import kd_loss -from . import lwf_loss -from . import dice_loss -from . import eqlv2 -from . import focal_loss -from . import softiou_loss -from . import dmt_loss -from . import reco_loss -from . import reverse_loss \ No newline at end of file diff --git a/mpa/modules/experimental/models/losses/dice_loss.py b/mpa/modules/experimental/models/losses/dice_loss.py deleted file mode 100644 index 91314614..00000000 --- a/mpa/modules/experimental/models/losses/dice_loss.py +++ /dev/null @@ -1,165 +0,0 @@ -"""Modified from https://github.com/LikeLy-Journey/SegmenTron/blob/master/ -segmentron/solver/loss.py (Apache-2.0 License)""" -import torch -import torch.nn as nn -import torch.nn.functional as F - -# from ..builder import LOSSES -from mmseg.models.builder import LOSSES -from mpa.modules.models.losses.utils import weighted_loss -from mpa.modules.models.losses.utils_pixel_wise import builder - - -@weighted_loss -def dice_loss(pred, - target, - valid_mask, - smooth=1, - exponent=2, - class_weight=None, - ignore_index=255): - assert pred.shape[0] == target.shape[0] - total_loss = 0 - num_classes = pred.shape[1] - for i in range(num_classes): - if i != ignore_index: - dice_loss = binary_dice_loss( - pred[:, i], - target[..., i], - valid_mask=valid_mask, - smooth=smooth, - exponent=exponent) - if class_weight is not None: - dice_loss *= class_weight[i] - total_loss += dice_loss - return total_loss / num_classes - - -@weighted_loss -def binary_dice_loss(pred, target, valid_mask, smooth=1, exponent=2, **kwargs): - assert pred.shape[0] == target.shape[0] - pred = pred.reshape(pred.shape[0], -1) - target = target.reshape(target.shape[0], -1) - valid_mask = valid_mask.reshape(valid_mask.shape[0], -1) - - num = torch.sum(torch.mul(pred, target) * valid_mask, dim=1) * 2 + smooth - den = torch.sum(pred.pow(exponent) + target.pow(exponent), dim=1) + smooth - - return 1 - num / den - - -@LOSSES.register_module() -class DiceLossMPA(nn.Module): - """DiceLoss. - - This loss is proposed in `V-Net: Fully Convolutional Neural Networks for - Volumetric Medical Image Segmentation `_. - - Args: - loss_type (str, optional): Binary or multi-class loss. - Default: 'multi_class'. Options are "binary" and "multi_class". - smooth (float): A float number to smooth loss, and avoid NaN error. - Default: 1 - exponent (float): An float number to calculate denominator - value: \\sum{x^exponent} + \\sum{y^exponent}. Default: 2. - reduction (str, optional): The method used to reduce the loss. Options - are "none", "mean" and "sum". This parameter only works when - per_image is True. Default: 'mean'. - class_weight (list[float], optional): The weight for each class. - Default: None. - loss_weight (float, optional): Weight of the loss. Default to 1.0. - ignore_index (int | None): The label index to be ignored. Default: 255. - """ - - def __init__(self, - smooth=1, - exponent=2, - reduction='mean', - class_weight=None, - loss_weight=1.0, - ignore_index=255, - **kwargs): - super(DiceLossMPA, self).__init__() - self.smooth = smooth - self.exponent = exponent - self.reduction = reduction - self.class_weight = class_weight - self.loss_weight = loss_weight - self.ignore_index = ignore_index - - # for pixel-based segmentation loss - self._loss_weight_scheduler = builder.build_scheduler(loss_weight, default_value=1.0) - - self._iter = 0 - self._last_loss_weight = 0 - - from mpa.utils.logger import get_logger - logger = get_logger() - logger.info('################################ TODO #################################') - logger.info('# Some meta cfgs used in BasePixelLoss is not usable in DiceLoss yet. #') - logger.info('# (reg_weight, scale, raw_sparsity, weight_sparsity) #') - logger.info('#######################################################################') - - @property - def iter(self): - return self._iter - - @property - def last_loss_weight(self): - return self._last_loss_weight - - @property - def name(self): - return 'dice' - - def _forward(self, - pred, - target, - avg_factor=None, - reduction_override=None, - **kwargs): - assert reduction_override in (None, 'none', 'mean', 'sum') - reduction = ( - reduction_override if reduction_override else self.reduction) - if self.class_weight is not None: - class_weight = pred.new_tensor(self.class_weight) - else: - class_weight = None - - pred = F.softmax(pred, dim=1) - num_classes = pred.shape[1] - one_hot_target = F.one_hot( - torch.clamp(target.long(), 0, num_classes - 1), - num_classes=num_classes) - valid_mask = (target != self.ignore_index).long() - - losses = self.loss_weight * dice_loss( - pred, - one_hot_target, - valid_mask=valid_mask, - reduction=reduction, - avg_factor=avg_factor, - smooth=self.smooth, - exponent=self.exponent, - class_weight=class_weight, - ignore_index=self.ignore_index) - - meta = dict( - weight=self.last_loss_weight, - # reg_weight=self.last_reg_weight, # TODO - # scale=self.last_scale, # TODO - # raw_sparsity=raw_sparsity, # TODO - # weight_sparsity=weight_sparsity # TODO - ) - - return losses, meta - - def forward(self, *args, **kwargs): - self._last_loss_weight = self._loss_weight_scheduler(self.iter) - - loss, meta = self._forward(*args, **kwargs) - out_loss = self._last_loss_weight * loss - - self._iter += 1 - - return out_loss, meta diff --git a/mpa/modules/experimental/models/losses/distillation_losses.py b/mpa/modules/experimental/models/losses/distillation_losses.py deleted file mode 100644 index 8426e353..00000000 --- a/mpa/modules/experimental/models/losses/distillation_losses.py +++ /dev/null @@ -1,55 +0,0 @@ -import torch -import torch.nn as nn -import torch.nn.functional as F - -from mmcls.models.builder import LOSSES -from mmcls.models.losses.utils import weight_reduce_loss - - -def lwf_loss(pred, label, T=2.0, weight=None, reduction='mean', avg_factor=None): - # # element-wise losses - # label = F.softmax(torch.from_numpy(label/T).cuda(), dim=1) - label = F.softmax(label / T, dim=1) - loss = -torch.sum(F.log_softmax(pred/T, 1) * label) - # apply weights and do the reduction - if weight is not None: - weight = weight.float() - loss = weight_reduce_loss( - loss, weight=weight, reduction=reduction, avg_factor=avg_factor) - - return loss - - -@LOSSES.register_module() -class LwfLoss(nn.Module): - - def __init__(self, - T=2.0, - reduction='mean', - loss_weight=1.0): - super(LwfLoss, self).__init__() - self.T = T - self.reduction = reduction - self.loss_weight = loss_weight - - self.criterion = lwf_loss - - def forward(self, - cls_score, - label, - weight=None, - avg_factor=None, - reduction_override=None, - **kwargs): - assert reduction_override in (None, 'none', 'mean', 'sum') - reduction = ( - reduction_override if reduction_override else self.reduction) - loss_cls = self.loss_weight * self.criterion( - cls_score, - label, - self.T, - weight, - reduction=reduction, - avg_factor=avg_factor, - **kwargs) - return loss_cls diff --git a/mpa/modules/experimental/models/losses/dmt_loss.py b/mpa/modules/experimental/models/losses/dmt_loss.py deleted file mode 100644 index e2125d3b..00000000 --- a/mpa/modules/experimental/models/losses/dmt_loss.py +++ /dev/null @@ -1,229 +0,0 @@ -import torch -import torch.nn.functional as F -from mmseg.models.builder import LOSSES -from mpa.modules.models.losses.utils import get_class_weight -from mpa.modules.models.losses.pixel_base import BasePixelLoss -from mmseg.models.losses.utils import weight_reduce_loss - - -def dmt_loss(pred, - label, - weight=None, - class_weight=None, - reduction='mean', - avg_factor=None, - ignore_index=255, - alpha=1.0, - gamma1=5, - gamma2=5): - """The wrapper function for :func:`F.cross_entropy`""" - # if labeled batch - if len(label.size()) <= 3: - loss = F.cross_entropy( - pred, - label, - weight=class_weight, - reduction='none', - ignore_index=255) - - # if unlabeled batch - else: - student_logit = pred - teacher_logit = label - student_softmax_logit = student_logit.softmax(dim=1).clone().detach() - student_pseudo_confidence, student_pseudo_label = torch.max(student_softmax_logit, dim=1) - - teacher_softmax_logit = teacher_logit.softmax(dim=1).clone().detach() - teacher_pseudo_confidence, teacher_pseudo_label = torch.max(teacher_softmax_logit, dim=1) - - total_loss = F.cross_entropy(student_logit, teacher_pseudo_label, reduction='none', ignore_index=255) - - probabilities_current = student_softmax_logit.gather(dim=1, index=teacher_pseudo_label.unsqueeze(1)).squeeze(1) - dynamic_weights = torch.ones_like(probabilities_current).float() - - # Prepare indices - disagreement = student_pseudo_label != teacher_pseudo_label - current_win = torch.gt(student_pseudo_confidence, teacher_pseudo_confidence) - - # Agree - indices = ~disagreement - dynamic_weights[indices] = probabilities_current[indices] ** gamma1 - - # Disagree (current model wins, do not learn!) - indices = disagreement * current_win - dynamic_weights[indices] = 0 - - # Disagree - indices = disagreement * ~current_win - dynamic_weights[indices] = probabilities_current[indices] ** gamma2 - - # pseudo label maksing lower than top alpha percent confidence - pseudo_label_mask = teacher_pseudo_confidence < torch.quantile(teacher_pseudo_confidence, 1-alpha) - dynamic_weights[pseudo_label_mask] = 0 - valid_label = teacher_pseudo_label[~pseudo_label_mask] - - # Weight loss - loss = (total_loss * dynamic_weights).sum() / (valid_label != 255).sum() - - if weight is not None: - weight = weight.float() - loss = weight_reduce_loss( - loss, weight=weight, reduction=reduction, avg_factor=avg_factor) - - return loss - - -def _expand_onehot_labels(labels, target_shape, ignore_index): - """Expand onehot labels to match the size of prediction.""" - bin_labels = labels.new_zeros(target_shape) - valid_mask = (labels >= 0) & (labels != ignore_index) - inds = torch.nonzero(valid_mask, as_tuple=True) - - if inds[0].numel() > 0: - if labels.dim() == 3: - bin_labels[inds[0], labels[valid_mask], inds[1], inds[2]] = 1 - else: - bin_labels[inds[0], labels[valid_mask]] = 1 - - return bin_labels - - -def binary_cross_entropy(pred, - label, - class_weight=None, - ignore_index=255): - """Calculate the binary CrossEntropy loss. - - Args: - pred (torch.Tensor): The prediction with shape (N, 1). - label (torch.Tensor): The learning label of the prediction. - weight (torch.Tensor, optional): Sample-wise loss weight. - reduction (str, optional): The method used to reduce the loss. - Options are "none", "mean" and "sum". - avg_factor (int, optional): Average factor that is used to average - the loss. Defaults to None. - class_weight (list[float], optional): The weight for each class. - ignore_index (int | None): The label index to be ignored. Default: 255 - - Returns: - torch.Tensor: The calculated loss - """ - - if pred.dim() != label.dim(): - assert (pred.dim() == 2 and label.dim() == 1) or ( - pred.dim() == 4 and label.dim() == 3), \ - 'Only pred shape [N, C], label shape [N] or pred shape [N, C, ' \ - 'H, W], label shape [N, H, W] are supported' - label = _expand_onehot_labels(label, pred.shape, ignore_index) - - loss = F.binary_cross_entropy_with_logits( - pred, - label.float(), - pos_weight=class_weight, - reduction='none' - ) - - return loss - - -def mask_cross_entropy(pred, - target, - label, - class_weight=None, - ignore_index=None): - """Calculate the CrossEntropy loss for masks. - - Args: - pred (torch.Tensor): The prediction with shape (N, C), C is the number - of classes. - target (torch.Tensor): The learning label of the prediction. - label (torch.Tensor): ``label`` indicates the class label of the mask' - corresponding object. This will be used to select the mask in the - of the class which the object belongs to when the mask prediction - if not class-agnostic. - reduction (str, optional): The method used to reduce the loss. - Options are "none", "mean" and "sum". - avg_factor (int, optional): Average factor that is used to average - the loss. Defaults to None. - class_weight (list[float], optional): The weight for each class. - ignore_index (None): Placeholder, to be consistent with other loss. - Default: None. - - Returns: - torch.Tensor: The calculated loss - """ - - assert ignore_index is None, 'BCE loss does not support ignore_index' - - num_rois = pred.size()[0] - inds = torch.arange(0, num_rois, dtype=torch.long, device=pred.device) - pred_slice = pred[inds, label].squeeze(1) - - loss = F.binary_cross_entropy_with_logits( - pred_slice, - target, - weight=class_weight, - reduction='mean' - )[None] - - return loss - - -@LOSSES.register_module() -class DynamicMutualLoss(BasePixelLoss): - """CrossEntropyLoss. - - Args: - use_sigmoid (bool, optional): Whether the prediction uses sigmoid - of softmax. Defaults to False. - use_mask (bool, optional): Whether to use mask cross entropy loss. - Defaults to False. - class_weight (list[float] | str, optional): Weight of each class. If in - str format, read them from a file. Defaults to None. - """ - - def __init__(self, - use_sigmoid=False, - use_mask=False, - class_weight=None, - alpha=1.0, - gamma1=5, - gamma2=5, - **kwargs): - super(DynamicMutualLoss, self).__init__(**kwargs) - - self.use_sigmoid = use_sigmoid - self.use_mask = use_mask - self.class_weight = get_class_weight(class_weight) - self.alpha = alpha - self.gamma1 = gamma1 - self.gamma2 = gamma2 - assert (use_sigmoid is False) or (use_mask is False) - if self.use_sigmoid: - self.cls_criterion = binary_cross_entropy - elif self.use_mask: - self.cls_criterion = mask_cross_entropy - else: - self.cls_criterion = dmt_loss - - @property - def name(self): - return 'dmt' - - def _calculate(self, cls_score, label, scale, weight=None): - class_weight = None - if self.class_weight is not None: - class_weight = cls_score.new_tensor(self.class_weight) - - loss = self.cls_criterion( - scale * cls_score, - label, - weight, - class_weight=class_weight, - ignore_index=self.ignore_index, - alpha=self.alpha, - gamma1=self.gamma1, - gamma2=self.gamma2 - ) - - return loss, cls_score diff --git a/mpa/modules/experimental/models/losses/eqlv2.py b/mpa/modules/experimental/models/losses/eqlv2.py deleted file mode 100644 index e12d76fa..00000000 --- a/mpa/modules/experimental/models/losses/eqlv2.py +++ /dev/null @@ -1,96 +0,0 @@ -import torch -import torch.nn.functional as F -import torch.distributed as dist -from functools import partial - -from mmseg.models.builder import LOSSES -from mpa.modules.models.losses.pixel_base import BasePixelLoss -from mmseg.models.losses.cross_entropy_loss import _expand_onehot_labels -from mpa.utils.logger import get_logger - -logger = get_logger() - - -@LOSSES.register_module() -class EQLv2(BasePixelLoss): - def __init__(self, - gamma=12, - mu=0.8, - alpha=4.0, - vis_grad=False, - **kwargs): - - super(EQLv2, self).__init__(**kwargs) - self.group = True - - # cfg for eqlv2 - self.vis_grad = vis_grad - self.gamma = gamma - self.mu = mu - self.alpha = alpha - - # initial variables - self._pos_grad = None - self._neg_grad = None - self.pos_neg = None - - def _func(x, gamma, mu): - return 1 / (1 + torch.exp(-gamma * (x - mu))) - self.map_func = partial(_func, gamma=self.gamma, mu=self.mu) - logger.info(f"build EQL v2, gamma: {gamma}, mu: {mu}, alpha: {alpha}") - - @property - def name(self): - return 'eqlv2' - - def _calculate(self, cls_score, label, scale, weight=None): - cls_score *= scale - self.n_i, self.n_c, self.h, self.w = cls_score.size() - - target = _expand_onehot_labels(label, cls_score.size(), self.ignore_index) - - pos_w, neg_w = self.get_weight(cls_score) - - weight = pos_w * target + neg_w * (1 - target) - - # (16, 21, 512, 512) -> (16, 21, 512, 512) - cls_loss = F.binary_cross_entropy_with_logits(cls_score, target.float(), reduction='none') - cls_loss = torch.mean(cls_loss * weight, dim=1) # (N, 512, 512) - - self.collect_grad(cls_score.detach(), target.detach(), weight.detach()) - - return cls_loss, cls_score - - def collect_grad(self, cls_score, target, weight): - prob = torch.sigmoid(cls_score) - - grad = target * (prob - 1) + (1 - target) * prob - grad = torch.abs(grad) - - pos_grad = torch.sum(grad * target * weight, dim=0) - neg_grad = torch.sum(grad * (1 - target) * weight, dim=0) - - if torch.cuda.device_count() > 1: - dist.all_reduce(pos_grad) - dist.all_reduce(neg_grad) - - self._pos_grad += pos_grad - self._neg_grad += neg_grad - - self.pos_neg = self._pos_grad.mean(dim=(1, 2)) / (self._neg_grad.mean(dim=(1, 2)) + 1e-10) # (21) - - def get_weight(self, cls_score): - # we do not have information about pos grad and neg grad at beginning - if self._pos_grad is None: - self._pos_grad = cls_score.new_zeros((self.n_c, self.h, self.w)) - self._neg_grad = cls_score.new_zeros((self.n_c, self.h, self.w)) - neg_w = cls_score.new_ones((self.n_i, self.n_c, self.h, self.w)) - pos_w = cls_score.new_ones((self.n_i, self.n_c, self.h, self.w)) - else: - neg_w = self.map_func(self.pos_neg) - pos_w = 1 + self.alpha * (1 - neg_w) - - neg_w = neg_w.view(1, self.n_c, 1, 1).expand(self.n_i, self.n_c, self.h, self.w) - pos_w = pos_w.view(1, self.n_c, 1, 1).expand(self.n_i, self.n_c, self.h, self.w) - - return pos_w, neg_w diff --git a/mpa/modules/experimental/models/losses/focal_loss.py b/mpa/modules/experimental/models/losses/focal_loss.py deleted file mode 100644 index 89e8992a..00000000 --- a/mpa/modules/experimental/models/losses/focal_loss.py +++ /dev/null @@ -1,54 +0,0 @@ -import torch.nn.functional as F - -from mmseg.models.builder import LOSSES -from mpa.modules.models.losses.pixel_base import BasePixelLoss -from mmseg.models.losses.cross_entropy_loss import _expand_onehot_labels - - -def focal_loss(pred, - target, - gamma=2.0, - alpha=0.25, - ignore_index=255): - - pred_softmax = F.softmax(pred, 1) - target = _expand_onehot_labels(target, pred.shape, ignore_index) - - pt = (1 - pred_softmax) * target + pred_softmax * (1 - target) - focal_weight = (alpha * target + (1 - alpha) * (1 - target)) * pt.pow(gamma) - loss = F.binary_cross_entropy_with_logits(pred, target.float(), reduction='none') - loss *= focal_weight - return loss.mean(dim=1) - - -@LOSSES.register_module() -class FocalLossMPA(BasePixelLoss): - def __init__(self, gamma=2.0, alpha=0.25, **kwargs): - super(FocalLossMPA, self).__init__(**kwargs) - - self.gamma = gamma - self.alpha = alpha - - self.cls_criterion = focal_loss - - from mpa.utils.logger import get_logger - logger = get_logger() - logger.info('######################################################################################') - logger.info('# This Focal loss is implemented based on SoftmaxFocalLoss at class_balanced_loss.py #') - logger.info('# Differently, label smoothing is not considered. #') - logger.info('######################################################################################') - - @property - def name(self): - return 'focal' - - def _calculate(self, cls_score, label, scale, weight=None): - loss = self.cls_criterion( - scale * cls_score, - label, - gamma=self.gamma, - alpha=self.alpha, - ignore_index=self.ignore_index - ) - - return loss, cls_score diff --git a/mpa/modules/experimental/models/losses/kd_loss.py b/mpa/modules/experimental/models/losses/kd_loss.py deleted file mode 100644 index d05cb14b..00000000 --- a/mpa/modules/experimental/models/losses/kd_loss.py +++ /dev/null @@ -1,67 +0,0 @@ -import torch -import torch.nn as nn -import torch.nn.functional as F - -from mmdet.models import LOSSES -from mmdet.models.losses import weight_reduce_loss - - -@LOSSES.register_module() -class KDLoss(nn.Module): - def __init__(self, - use_sigmoid=False, - temperature=2.0, - reduction='mean', - loss_weight=1.0): - """Knowledege Distillation loss - - Args: - use_sigmoid (bool, optional): Whether the prediction uses sigmoid - of softmax. Defaults to False. - reduction (str, optional): . Defaults to 'mean'. - Options are "none", "mean" and "sum". - loss_weight (float, optional): Weight of the loss. Defaults to 1.0. - """ - super().__init__() - self.use_sigmoid = use_sigmoid - self.T = temperature - self.reduction = reduction - self.loss_weight = loss_weight - - def forward(self, - input, - target, - weight=None, - avg_factor=None, - **kwargs): - """Forward function. - - Args: - input (torch.Tensor): The prediction (N, C, ...) - target (torch.Tensor): The learning target of the prediction (N, C, ...) - weight (torch.Tensor, optional): Sample-wise loss weight (N, 1, ...) - avg_factor (int, optional): Average factor that is used to average - the loss. Defaults to None. - Returns: - torch.Tensor: The calculated loss - """ - - # Elem-wise KD loss - if self.use_sigmoid: - input_prob = torch.sigmoid(input/self.T) - target_prob = torch.sigmoid(target/self.T) - losses = (target_prob - input_prob)**2 - else: - input_log_prob = F.log_softmax(input/self.T, dim=1) - target_prob = F.softmax(target/self.T, dim=1) - losses = (-target_prob*input_log_prob).sum(dim=1) - - # losses = F.kl_div(input_log_prob, target_prob, reduction='none') - - # Weighted reduction - if weight is not None: - weight = weight.float() - loss = weight_reduce_loss( - losses, weight=weight, reduction=self.reduction, avg_factor=avg_factor) - - return loss*self.loss_weight diff --git a/mpa/modules/experimental/models/losses/lwf_loss.py b/mpa/modules/experimental/models/losses/lwf_loss.py deleted file mode 100644 index c12274bb..00000000 --- a/mpa/modules/experimental/models/losses/lwf_loss.py +++ /dev/null @@ -1,191 +0,0 @@ -import torch -import torch.nn as nn -import torch.nn.functional as F - -from mmdet.models import LOSSES -from mmdet.models.losses import weight_reduce_loss -from .kd_loss import KDLoss -from mpa.modules.utils.task_adapt import map_class_names - - -@LOSSES.register_module() -class LwFLoss(KDLoss): - def __init__(self, - src_classes, - dst_classes, - bg_aware=True, - temperature=2.0, - reduction='mean', - loss_weight=1.0): - """ loss - - Args: - src_classes (list[str]): Teacher (OLD) classes - dst_classes (list[str]): Student (NEW) classes - bg_aware (bool, optional): Whether to enable BG-aware distillation - '__bg__' class would be added the end of src/dst_classes - temperature (float, optional): Temperature for KD loss - reduction (str, optional): . Defaults to 'mean'. - Options are "none", "mean" and "sum". - loss_weight (float, optional): Weight of the loss. Defaults to 1.0. - """ - print('LwFLoss init!') - super().__init__( - use_sigmoid=False, # softmax only - temperature=temperature, - reduction=reduction, - loss_weight=loss_weight - ) - self.src_classes = src_classes.copy() - self.dst_classes = dst_classes.copy() - self.bg_aware = bg_aware - if bg_aware: - self.src_classes += ['__bg__'] - self.dst_classes += ['__bg__'] - self.src2dst = torch.tensor(map_class_names(self.src_classes, self.dst_classes)) - self.dst2src = torch.tensor(map_class_names(self.dst_classes, self.src_classes)) - - def forward(self, - input, - target, - target_is_logit=True, - weight=None, - avg_factor=None, - **kwargs): - """Forward function. - - Args: - input (torch.Tensor): The prediction (N, C_new, ...) - target (torch.Tensor): The learning target of the prediction (N, C_old, ...) - weight (torch.Tensor, optional): Sample-wise loss weight (N, 1, ...) - avg_factor (int, optional): Average factor that is used to average - the loss. Defaults to None. - Returns: - torch.Tensor: The calculated loss - """ - - if input.shape[0] == 0: - return torch.tensor(0.0) - - if not self.bg_aware: - # Gathering predictions according to class name matching - src_pred = target[:, self.src2dst >= 0] - dst_pred = input[:, self.src2dst[self.src2dst >= 0]] - return super().forward(dst_pred, src_pred, weight=weight, avg_factor=avg_factor, **kwargs) - - # --- Elem-wise BG-aware LwF loss - # Non-matching probabilities would be regarded as 'BACKGROUND' probability - # So, they are accumulated to BG probability (both for OLD & NEW) - if target_is_logit: - src_prob_all = F.softmax(target/self.T, dim=1) - else: - src_prob_all = target - src_prob_gathered = src_prob_all[:, self.src2dst >= 0] - src_prob_nomatch = src_prob_all[:, self.src2dst < 0] - src_prob_gathered[:, -1] += src_prob_nomatch.sum(dim=1) # Accumulating to BG prob - dst_prob_all = F.softmax(input/self.T, dim=1) - dst_prob_gathered = dst_prob_all[:, self.src2dst[self.src2dst >= 0]] - dst_prob_nomatch = dst_prob_all[:, self.dst2src < 0] - dst_prob_gathered[:, -1] += dst_prob_nomatch.sum(dim=1) # Accumulating to BG prob - - # src_logit_gathered = target[:, self.src2dst >= 0] - # src_prob_gathered = F.softmax(src_logit_gathered/self.T, dim=1) - # dst_logit_gathered = input[:, self.src2dst[self.src2dst >= 0]] - # dst_prob_gathered = F.softmax(dst_logit_gathered/self.T, dim=1) - - # X-entropy - losses = -src_prob_gathered*torch.log(dst_prob_gathered) - # losses = losses.sum(dim=1) - - # --- Weighted reduction - if weight is not None: - weight = weight.float() - loss = weight_reduce_loss( - losses, weight=weight, reduction=self.reduction, avg_factor=avg_factor) - - return loss*self.loss_weight - - -@LOSSES.register_module() -class ClassSensitiveCrossEntropyLoss(nn.Module): - def __init__(self, - model_classes, - data_classes, - bg_aware=True, - reduction='mean', - loss_weight=1.0): - """ loss - - Args: - model_classes (list[str]): Model classes - data_classes (list[str]): Data classes - bg_aware (bool, optional): Whether to enable BG-aware distillation - '__bg__' class would be added the end of model/data_classes - temperature (float, optional): Temperature for KD loss - reduction (str, optional): . Defaults to 'mean'. - Options are "none", "mean" and "sum". - loss_weight (float, optional): Weight of the loss. Defaults to 1.0. - """ - print('ClassSensitiveCrossEntropyLoss init!') - super().__init__() - self.model_classes = model_classes.copy() - self.data_classes = data_classes.copy() - self.bg_aware = bg_aware - self.reduction = reduction - self.loss_weight = loss_weight - if bg_aware: - self.model_classes += ['__bg__'] - self.data_classes += ['__bg__'] - self.model2data = torch.tensor(map_class_names(self.model_classes, self.data_classes)) - # print(f'model classes = {self.model_classes}, data classes = {self.data_classes}') - # print(f'########### model2data mapping = {self.model2data}') - - def forward(self, - logits, - labels, - weight=None, - avg_factor=None, - **kwargs): - """Forward function. - - Args: - logits (torch.Tensor): The prediction (N, C_new, ...) - labels (torch.Tensor): The learning target of the prediction (N, ...) - weight (torch.Tensor, optional): Sample-wise loss weight (N, 1, ...) - avg_factor (int, optional): Average factor that is used to average - the loss. Defaults to None. - Returns: - torch.Tensor: The calculated loss - """ - - if logits.shape[0] == 0: - return torch.tensor(0.0) - - # --- Elem-wise BG-aware X-entropy loss - # Non-matching probabilities would be regarded as 'BACKGROUND' probability - # So, they are accumulated to BG probability (both for input & target) - prob_all = F.softmax(logits, dim=1) - # print(f'shape of prob_all {prob_all.shape}') - # print(f'shape of prob_all model2data {self.model2data.shape}') - prob_gathered = prob_all[:, self.model2data >= 0] - prob_nomatch = prob_all[:, self.model2data < 0] - prob_gathered[:, -1] += prob_nomatch.sum(dim=1) # Accumulating non-matching probs to BG prob - prob_log = torch.log(prob_gathered) - label_gathered = self.model2data[labels].to(labels) - - # X-entropy: NLL loss w/ log(softmax(logit)) & labels - losses = F.nll_loss(prob_log, label_gathered, reduction='none') - - # logit_gathered = logits[:, self.model2data >= 0] - # label_gathered = self.model2data[labels].to(labels) - # logit_gathered = logits - # label_gathered = labels - # losses = F.cross_entropy(logit_gathered, label_gathered, reduction='none') - - # --- Weighted reduction - if weight is not None: - weight = weight.float() - loss = weight_reduce_loss( - losses, weight=weight, reduction=self.reduction, avg_factor=avg_factor) - - return loss*self.loss_weight diff --git a/mpa/modules/experimental/models/losses/reco_loss.py b/mpa/modules/experimental/models/losses/reco_loss.py deleted file mode 100644 index 9f0842e2..00000000 --- a/mpa/modules/experimental/models/losses/reco_loss.py +++ /dev/null @@ -1,141 +0,0 @@ -import torch -import torch.nn as nn -import torch.nn.functional as F -import numpy as np - -from mmseg.models.builder import LOSSES - - -def label_onehot(inputs, num_segments): - batch_size, _, im_h, im_w = inputs.shape - # remap invalid pixels (-1) into 0, otherwise we cannot create one-hot vector with negative labels. - # we will still mask out those invalid values in valid mask - inputs = torch.relu(inputs) - outputs = torch.zeros([batch_size, num_segments, im_h, im_w]).to(inputs.device) - return outputs.scatter_(1, inputs.long(), 1.0) - - -def negative_index_sampler(samp_num, seg_num_list): - negative_index = [] - for i in range(samp_num.shape[0]): - for j in range(samp_num.shape[1]): - sum1 = sum(seg_num_list[:j]) - sum2 = sum(seg_num_list[:j+1]) - if sum1 > sum2: - tmp = sum1 - sum1 = sum2 - sum2 = tmp - print("Warning : ValueError: low >= high") - negative_index += np.random.randint(low=sum1, - high=sum2, - size=int(samp_num[i, j])).tolist() - return negative_index - - -def compute_reco_loss(rep, label, mask, prob, strong_threshold=1.0, temp=0.5, num_queries=256, num_negatives=256): - # batch_size, num_feat, im_w_, im_h = rep.shape - _, num_feat, _, _ = rep.shape - num_segments = prob.shape[1] - device = rep.device - - # compute valid binary mask for each pixel - # when we use cutmix(GT+unlabeled), GT may contain the ignore label(255). - # However, since 255 cannot be converted via label_onehot(), it makes error. - ignore_mask = (label > num_segments-1) - label[ignore_mask] = 0 - mask[ignore_mask] = 0.0 - label1 = label_onehot(label, num_segments) - valid_pixel = label1 * mask - - # permute representation for indexing: batch x im_h x im_w x feature_channel - rep = rep.permute(0, 2, 3, 1) - - # compute prototype (class mean representation) for each class across all valid pixels - seg_feat_all_list = [] - seg_feat_hard_list = [] - seg_num_list = [] - seg_proto_list = [] - for i in range(num_segments): - valid_pixel_seg = valid_pixel[:, i] # select binary mask for i-th class - if valid_pixel_seg.sum() == 0: # not all classes would be available in a mini-batch - continue - - prob_seg = prob[:, i, :, :] - rep_mask_hard = (prob_seg < strong_threshold) * valid_pixel_seg.bool() # select hard queries - - seg_proto_list.append(torch.mean(rep[valid_pixel_seg.bool()], dim=0, keepdim=True)) - seg_feat_all_list.append(rep[valid_pixel_seg.bool()]) - seg_feat_hard_list.append(rep[rep_mask_hard]) - seg_num_list.append(int(valid_pixel_seg.sum().item())) - - # compute regional contrastive loss - if len(seg_num_list) <= 1: # in some rare cases, a small mini-batch might only contain 1 or no semantic class - return torch.tensor(0.0) - else: - reco_loss = torch.tensor(0.0) - seg_proto = torch.cat(seg_proto_list) - valid_seg = len(seg_num_list) - seg_len = torch.arange(valid_seg) - - for i in range(valid_seg): - # sample hard queries - if len(seg_feat_hard_list[i]) > 0: - seg_hard_idx = torch.randint(len(seg_feat_hard_list[i]), size=(num_queries,)) - anchor_feat_hard = seg_feat_hard_list[i][seg_hard_idx] - anchor_feat = anchor_feat_hard - else: # in some rare cases, all queries in the current query class are easy - continue - - # apply negative key sampling (with no gradients) - with torch.no_grad(): - # generate index mask for the current query class; e.g. [0, 1, 2] -> [1, 2, 0] -> [2, 0, 1] - seg_mask = torch.cat(([seg_len[i:], seg_len[:i]])) - - # compute similarity for each negative segment prototype (semantic class relation graph) - proto_sim = torch.cosine_similarity(seg_proto[seg_mask[0]].unsqueeze(0), seg_proto[seg_mask[1:]], dim=1) - proto_prob = torch.softmax(proto_sim / temp, dim=0) - - # sampling negative keys based on the generated distribution [num_queries x num_negatives] - negative_dist = torch.distributions.categorical.Categorical(probs=proto_prob) - samp_class = negative_dist.sample(sample_shape=[num_queries, num_negatives]) - samp_num = torch.stack([(samp_class == c).sum(1) for c in range(len(proto_prob))], dim=1) - - # sample negative indices from each negative class - negative_num_list = seg_num_list[i+1:] + seg_num_list[:i] - negative_index = negative_index_sampler(samp_num, negative_num_list) - - # index negative keys (from other classes) - negative_feat_all = torch.cat(seg_feat_all_list[i+1:] + seg_feat_all_list[:i]) - negative_feat = negative_feat_all[negative_index].reshape(num_queries, num_negatives, num_feat) - - # combine positive and negative keys: keys = [positive key | negative keys] with 1 + num_negative dim - positive_feat = seg_proto[i].unsqueeze(0).unsqueeze(0).repeat(num_queries, 1, 1) - all_feat = torch.cat((positive_feat, negative_feat), dim=1) - - seg_logits = torch.cosine_similarity(anchor_feat.unsqueeze(1), all_feat, dim=2) - reco_loss = reco_loss + F.cross_entropy(seg_logits / temp, torch.zeros(num_queries).long().to(device)) - return reco_loss / valid_seg - - -@LOSSES.register_module() -class ReCoLoss(nn.Module): - - def __init__(self, reduction='mean', loss_weight=1.0, ignore_index=None): - super(ReCoLoss, self).__init__() - self.reduction = reduction - self.loss_weight = loss_weight - self.ignore_index = ignore_index - - self.cls_criterion = compute_reco_loss - - def forward(self, - rep, label, mask, prob, - strong_threshold=1.0, temp=0.5, - num_queries=256, num_negatives=256, - **kwargs): - - loss_cls = self.loss_weight * self.cls_criterion(rep, label, mask, prob, - strong_threshold=strong_threshold, temp=temp, - num_queries=num_queries, num_negatives=num_negatives) - - return loss_cls diff --git a/mpa/modules/experimental/models/losses/reverse_loss.py b/mpa/modules/experimental/models/losses/reverse_loss.py deleted file mode 100644 index 53a747c7..00000000 --- a/mpa/modules/experimental/models/losses/reverse_loss.py +++ /dev/null @@ -1,193 +0,0 @@ -import torch -import torch.nn.functional as F -from mmseg.models.builder import LOSSES -from mpa.modules.models.losses.utils import get_class_weight -from mpa.modules.models.losses.pixel_base import BasePixelLoss - - -def reverse_loss(pred, - label, - weight=None, - class_weight=None, - reduction='mean', - avg_factor=None, - ignore_index=255): - """The wrapper function for :func:`F.cross_entropy`""" - - if len(label.size()) <= 3: - loss = F.cross_entropy( - pred, - label, - weight=class_weight, - reduction='none', - ignore_index=255) - # print('sup!') - else: - student_logit = pred - teacher_logit = label - - student_softmax_logit = student_logit.softmax(dim=1).clone().detach() - student_pseudo_confidence, student_pseudo_label = torch.max(student_softmax_logit, dim=1) - - teacher_softmax_logit = teacher_logit.softmax(dim=1).clone().detach() - teacher_pseudo_confidence, teacher_pseudo_label = torch.max(teacher_softmax_logit, dim=1) - - # self correction loss is only applied to pseudo_labels. - probabilities_current = student_softmax_logit.gather(dim=1, index=teacher_pseudo_label.unsqueeze(1)).squeeze(1) - w = torch.max(student_softmax_logit, dim=1)[0] - - first = (-1*(w * teacher_pseudo_confidence.to(dtype=w.dtype) * torch.log(probabilities_current))) - second = (-1*(1-w) * probabilities_current * torch.log(teacher_pseudo_confidence.to(dtype=w.dtype))) - - loss = torch.mean(first) + torch.mean(second) - # print('unsup!') - - return loss - - -def _expand_onehot_labels(labels, target_shape, ignore_index): - """Expand onehot labels to match the size of prediction.""" - bin_labels = labels.new_zeros(target_shape) - valid_mask = (labels >= 0) & (labels != ignore_index) - inds = torch.nonzero(valid_mask, as_tuple=True) - - if inds[0].numel() > 0: - if labels.dim() == 3: - bin_labels[inds[0], labels[valid_mask], inds[1], inds[2]] = 1 - else: - bin_labels[inds[0], labels[valid_mask]] = 1 - - return bin_labels - - -def binary_cross_entropy(pred, - label, - class_weight=None, - ignore_index=255): - """Calculate the binary CrossEntropy loss. - - Args: - pred (torch.Tensor): The prediction with shape (N, 1). - label (torch.Tensor): The learning label of the prediction. - weight (torch.Tensor, optional): Sample-wise loss weight. - reduction (str, optional): The method used to reduce the loss. - Options are "none", "mean" and "sum". - avg_factor (int, optional): Average factor that is used to average - the loss. Defaults to None. - class_weight (list[float], optional): The weight for each class. - ignore_index (int | None): The label index to be ignored. Default: 255 - - Returns: - torch.Tensor: The calculated loss - """ - - if pred.dim() != label.dim(): - assert (pred.dim() == 2 and label.dim() == 1) or ( - pred.dim() == 4 and label.dim() == 3), \ - 'Only pred shape [N, C], label shape [N] or pred shape [N, C, ' \ - 'H, W], label shape [N, H, W] are supported' - label = _expand_onehot_labels(label, pred.shape, ignore_index) - - loss = F.binary_cross_entropy_with_logits( - pred, - label.float(), - pos_weight=class_weight, - reduction='none' - ) - - return loss - - -def mask_cross_entropy(pred, - target, - label, - class_weight=None, - ignore_index=None): - """Calculate the CrossEntropy loss for masks. - - Args: - pred (torch.Tensor): The prediction with shape (N, C), C is the number - of classes. - target (torch.Tensor): The learning label of the prediction. - label (torch.Tensor): ``label`` indicates the class label of the mask' - corresponding object. This will be used to select the mask in the - of the class which the object belongs to when the mask prediction - if not class-agnostic. - reduction (str, optional): The method used to reduce the loss. - Options are "none", "mean" and "sum". - avg_factor (int, optional): Average factor that is used to average - the loss. Defaults to None. - class_weight (list[float], optional): The weight for each class. - ignore_index (None): Placeholder, to be consistent with other loss. - Default: None. - - Returns: - torch.Tensor: The calculated loss - """ - - assert ignore_index is None, 'BCE loss does not support ignore_index' - - num_rois = pred.size()[0] - inds = torch.arange(0, num_rois, dtype=torch.long, device=pred.device) - pred_slice = pred[inds, label].squeeze(1) - - loss = F.binary_cross_entropy_with_logits( - pred_slice, - target, - weight=class_weight, - reduction='mean' - )[None] - - return loss - - -@LOSSES.register_module() -class ReverseLearningLoss(BasePixelLoss): - """CrossEntropyLoss. - - Args: - use_sigmoid (bool, optional): Whether the prediction uses sigmoid - of softmax. Defaults to False. - use_mask (bool, optional): Whether to use mask cross entropy loss. - Defaults to False. - class_weight (list[float] | str, optional): Weight of each class. If in - str format, read them from a file. Defaults to None. - """ - - def __init__(self, - use_sigmoid=False, - use_mask=False, - class_weight=None, - **kwargs): - super(ReverseLearningLoss, self).__init__(**kwargs) - - self.use_sigmoid = use_sigmoid - self.use_mask = use_mask - self.class_weight = get_class_weight(class_weight) - - assert (use_sigmoid is False) or (use_mask is False) - if self.use_sigmoid: - self.cls_criterion = binary_cross_entropy - elif self.use_mask: - self.cls_criterion = mask_cross_entropy - else: - self.cls_criterion = reverse_loss - - @property - def name(self): - return 'rev' - - def _calculate(self, cls_score, label, scale, weight=None): - class_weight = None - if self.class_weight is not None: - class_weight = cls_score.new_tensor(self.class_weight) - - loss = self.cls_criterion( - scale * cls_score, - label, - weight, - class_weight=class_weight, - ignore_index=self.ignore_index - ) - - return loss, cls_score diff --git a/mpa/modules/experimental/models/losses/softiou_loss.py b/mpa/modules/experimental/models/losses/softiou_loss.py deleted file mode 100644 index 16ab176c..00000000 --- a/mpa/modules/experimental/models/losses/softiou_loss.py +++ /dev/null @@ -1,140 +0,0 @@ -"""Modified from https://github.com/LikeLy-Journey/SegmenTron/blob/master/ -segmentron/solver/loss.py (Apache-2.0 License)""" -import torch.nn as nn -import torch.nn.functional as F - -# from ..builder import LOSSES -from mmseg.models.builder import LOSSES -from mpa.modules.models.losses.utils import weighted_loss -from mpa.modules.models.losses.utils_pixel_wise import builder -from mmseg.models.losses.cross_entropy_loss import _expand_onehot_labels - - -@weighted_loss -def softiou_loss(pred, - target, - valid_mask, - class_weight=None, - ignore_index=255): - - assert pred.shape[0] == target.shape[0] - - N, C, _, _ = pred.size() - - inter = pred * target - union = pred + target - inter - - inter = inter.view(N, C, -1).sum(dim=2) - union = union.view(N, C, -1).sum(dim=2) - - loss = inter / (union + 1e-16) - - return 1 - loss.mean() - - -@LOSSES.register_module() -class SoftIoULoss(nn.Module): - """DiceLoss. - - This loss is proposed in `V-Net: Fully Convolutional Neural Networks for - Volumetric Medical Image Segmentation `_. - - Args: - loss_type (str, optional): Binary or multi-class loss. - Default: 'multi_class'. Options are "binary" and "multi_class". - smooth (float): A float number to smooth loss, and avoid NaN error. - Default: 1 - exponent (float): An float number to calculate denominator - value: \\sum{x^exponent} + \\sum{y^exponent}. Default: 2. - reduction (str, optional): The method used to reduce the loss. Options - are "none", "mean" and "sum". This parameter only works when - per_image is True. Default: 'mean'. - class_weight (list[float], optional): The weight for each class. - Default: None. - loss_weight (float, optional): Weight of the loss. Default to 1.0. - ignore_index (int | None): The label index to be ignored. Default: 255. - """ - - def __init__(self, - reduction='mean', - class_weight=None, - loss_weight=1.0, - ignore_index=255, - **kwargs): - super(SoftIoULoss, self).__init__() - self.reduction = reduction - self.class_weight = class_weight - self.loss_weight = loss_weight - self.ignore_index = ignore_index - - # for pixel-based segmentation loss - self._loss_weight_scheduler = builder.build_scheduler(loss_weight, default_value=1.0) - - self._iter = 0 - self._last_loss_weight = 0 - - from mpa.utils.logger import get_logger - logger = get_logger() - logger.info('################################ TODO ####################################') - logger.info('# Some meta cfgs used in BasePixelLoss is not usable in SoftIoULoss yet. #') - logger.info('# (reg_weight, scale, raw_sparsity, weight_sparsity) #') - logger.info('##########################################################################') - - @property - def iter(self): - return self._iter - - @property - def last_loss_weight(self): - return self._last_loss_weight - - @property - def name(self): - return 'softiou' - - def _forward(self, - pred, - target, - avg_factor=None, - reduction_override=None, - **kwargs): - assert reduction_override in (None, 'none', 'mean', 'sum') - reduction = ( - reduction_override if reduction_override else self.reduction) - if self.class_weight is not None: - class_weight = pred.new_tensor(self.class_weight) - else: - class_weight = None - - pred = F.softmax(pred, dim=1) - one_hot_target = _expand_onehot_labels(target, pred.shape, self.ignore_index) - valid_mask = (target != self.ignore_index).long() - - losses = self.loss_weight * softiou_loss( - pred, - one_hot_target, - valid_mask=valid_mask, - reduction=reduction, - avg_factor=avg_factor, - class_weight=class_weight, - ignore_index=self.ignore_index) - - meta = dict( - weight=self.last_loss_weight, - # reg_weight=self.last_reg_weight, # TODO - # scale=self.last_scale, # TODO - # raw_sparsity=raw_sparsity, # TODO - # weight_sparsity=weight_sparsity # TODO - ) - - return losses, meta - - def forward(self, *args, **kwargs): - self._last_loss_weight = self._loss_weight_scheduler(self.iter) - - loss, meta = self._forward(*args, **kwargs) - out_loss = self._last_loss_weight * loss - - self._iter += 1 - - return out_loss, meta diff --git a/mpa/modules/experimental/models/segmentors/__init__.py b/mpa/modules/experimental/models/segmentors/__init__.py deleted file mode 100644 index 90fe5df5..00000000 --- a/mpa/modules/experimental/models/segmentors/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -# flake8: noqa -from . import cutmix_seg_gt_naive -from . import cutmix_seg_naive -from . import mean_teacher_naive -from . import semisl_segmentor -from . import oversampling -from . import keep_output \ No newline at end of file diff --git a/mpa/modules/experimental/models/segmentors/cutmix/gt_utils.py b/mpa/modules/experimental/models/segmentors/cutmix/gt_utils.py deleted file mode 100644 index a21b9cf0..00000000 --- a/mpa/modules/experimental/models/segmentors/cutmix/gt_utils.py +++ /dev/null @@ -1,168 +0,0 @@ -import numpy as np -import random -import torch -from skimage.measure import label, regionprops - - -def padding_bbox_new(rectangles, size): - area = 0.5 * (size ** 2) - y0, x0, y1, x1 = rectangles - h = y1 - y0 - w = x1 - x0 - # upper_h = min(int(area/w), size) - # upper_w = min(int(area/h), size) - new_h = int(size*(np.exp(np.random.uniform(low=0.0, high=1.0, size=(1)) * np.log(0.5)))) - new_w = int(area/new_h) - delta_h = new_h - h - delta_w = new_w - w - y_ratio = y0/(size-y1+1) - x_ratio = x0/(size-x1+1) - x1 = min(x1+int(delta_w*(1/(1+x_ratio))), size) - x0 = max(x0-int(delta_w*(x_ratio/(1+x_ratio))), 0) - y1 = min(y1+int(delta_h*(1/(1+y_ratio))), size) - y0 = max(y0-int(delta_h*(y_ratio/(1+y_ratio))), 0) - return [y0, x0, y1, x1] - - -def sliming_bbox(rectangles, size): - area = 0.5 * (size ** 2) - y0, x0, y1, x1 = rectangles - h = y1 - y0 - w = x1 - x0 - lower_h = int(area/w) - if lower_h > h: - # print('wrong') - new_h = h - else: - new_h = random.randint(lower_h, h) - new_w = int(area/new_h) - if new_w > w: - # print('wrong') - new_w = w - 1 - delta_h = h - new_h - delta_w = w - new_w - prob = random.random() - if prob > 0.5: - y1 = max(random.randint(y1 - delta_h, y1), y0) - y0 = max(y1 - new_h, y0) - else: - y0 = min(random.randint(y0, y0 + delta_h), y1) - y1 = min(y0 + new_h, y1) - prob = random.random() - if prob > 0.5: - x1 = max(random.randint(x1 - delta_w, x1), x0) - x0 = max(x1 - new_w, x0) - else: - x0 = min(random.randint(x0, x0 + delta_w), x1) - x1 = min(x0 + new_w, x1) - return [y0, x0, y1, x1] - - -def init_cutmix(crop_size): - h = crop_size - w = crop_size - n_masks = 1 - prop_range_from = 0.15 - prop_range = 0.4 - mask_props = np.random.uniform(prop_range_from, prop_range, size=(n_masks, 1)) - y_props = np.exp(np.random.uniform(low=0.0, high=1.0, size=(n_masks, 1)) * np.log(mask_props)) - x_props = mask_props / y_props - sizes = np.round(np.stack([y_props, x_props], axis=2) * np.array((h, w))[None, None, :]) - positions = np.round((np.array((h, w))-sizes) * np.random.uniform(low=0.0, high=1.0, size=sizes.shape)) - rectangles = np.append(positions, positions+sizes, axis=2)[0, 0] - - return rectangles - - -def generate_cutmix_mask(pred, sample_cat, area_thresh=0.0001, no_pad=False, no_slim=False): - h, w = pred.shape[2], pred.shape[3] - valid_masks = np.zeros_like(pred) - n_boxes = 1 - for each_image_in_batch in range(pred.shape[0]): - valid_mask = np.zeros((h, w)) - values = list(np.unique(pred[each_image_in_batch][0])) - for i in range(n_boxes): - if not sample_cat[each_image_in_batch] in values or i > 0: - rectangles = init_cutmix(h) - y0, x0, y1, x1 = rectangles - valid_mask[int(y0):int(y1), int(x0):int(x1)] = 1 - valid_mask[int(y0):int(y1), int(x0):int(x1)] - else: - rectangles = generate_cutmix(pred[each_image_in_batch][0], sample_cat[each_image_in_batch], - area_thresh, no_pad=no_pad, no_slim=no_slim) - y0, x0, y1, x1 = rectangles - tmp = [[int(y0), int(y1/2), int(x0), int(x1/2)], - [int(y0), int(y1/2), int(x1/2), int(x1)], - [int(y1/2), int(y1), int(x0), int(x1/2)], - [int(y1/2), int(y1), int(x1/2), int(x1)]] - # tmp_cnt = 0 - for i in range(4): - if np.random.rand() < 0.5: - # tmp_cnt += 1 - val_mask = valid_mask[tmp[i][0]:tmp[i][1], tmp[i][2]:tmp[i][3]] - valid_mask[tmp[i][0]:tmp[i][1], tmp[i][2]:tmp[i][3]] = 1 - val_mask - - # valid_mask[int(y0):int(y1), int(x0):int(x1)] = 1 - valid_masks[each_image_in_batch][0] = torch.from_numpy(valid_mask).long() - - # if pred.is_cuda: - # valid_masks = valid_masks.cuda() - - return torch.Tensor(valid_masks) - - -def generate_cutmix(pred, cat, area_thresh, no_pad=False, no_slim=False): - h = pred.shape[0] - # print('h',h) - area_all = h ** 2 - pred = (pred == cat) * 1 - pred = label(pred) - prop = regionprops(pred) - values = np.unique(pred)[1:] - random.shuffle(values) - - flag = 0 - for value in values: - if np.sum(pred == value) > area_thresh*area_all: - flag = 1 - break - if flag == 1: - rectangles = prop[value-1].bbox - # area = prop[value-1].area - area = (rectangles[2]-rectangles[0])*(rectangles[3]-rectangles[1]) - if area >= 0.5*area_all and not no_slim: - rectangles = sliming_bbox(rectangles, h) - elif area < 0.5*area_all and not no_pad: - rectangles = padding_bbox_new(rectangles, h) - else: - pass - else: - rectangles = init_cutmix(h) - return rectangles - - -def update_cutmix_bank(cutmix_bank, preds_teacher_unsup, img_id, sample_id, area_thresh=0.0001): - # cutmix_bank [num_classes, len(dataset)] - # preds_teacher_unsup [2,num_classes,h,w] - area_all = preds_teacher_unsup.shape[-1]**2 - pred1 = preds_teacher_unsup[0].max(0)[1] # (h,w) - pred2 = preds_teacher_unsup[1].max(0)[1] # (h,w) - values1 = torch.unique(pred1) - values2 = torch.unique(pred2) - # for img1 - for idx in range(cutmix_bank.shape[0]): - if idx not in values1: - cutmix_bank[idx][img_id] = 0 - elif torch.sum(pred1 == idx) < area_thresh * area_all: - cutmix_bank[idx][img_id] = 0 - else: - cutmix_bank[idx][img_id] = 1 - # for img2 - for idx in range(cutmix_bank.shape[0]): - if idx not in values2: - cutmix_bank[idx][sample_id] = 0 - elif torch.sum(pred2 == idx) < area_thresh * area_all: - cutmix_bank[idx][sample_id] = 0 - else: - cutmix_bank[idx][sample_id] = 1 - - return cutmix_bank diff --git a/mpa/modules/experimental/models/segmentors/cutmix/mask_gen.py b/mpa/modules/experimental/models/segmentors/cutmix/mask_gen.py deleted file mode 100644 index 4054444e..00000000 --- a/mpa/modules/experimental/models/segmentors/cutmix/mask_gen.py +++ /dev/null @@ -1,74 +0,0 @@ -import numpy as np - - -class BoxMaskGenerator (object): - def __init__(self, prop_range=(0.25, 0.5), n_boxes=3, random_aspect_ratio=True, - prop_by_area=True, within_bounds=True, invert=True): - if isinstance(prop_range, float): - prop_range = (prop_range, prop_range) - self.prop_range = prop_range - self.n_boxes = n_boxes - self.random_aspect_ratio = random_aspect_ratio - self.prop_by_area = prop_by_area - self.within_bounds = within_bounds - self.invert = invert - - def generate_params(self, n_masks, mask_shape, rng=None): - """ - Box masks can be generated quickly on the CPU so do it there. - >>> boxmix_gen = BoxMaskGenerator((0.25, 0.25)) - >>> params = boxmix_gen.generate_params(256, (32, 32)) - >>> t_masks = boxmix_gen.torch_masks_from_params(params, (32, 32), 'cuda:0') - :param n_masks: number of masks to generate (batch size) - :param mask_shape: Mask shape as a `(height, width)` tuple - :param rng: [optional] np.random.RandomState instance - :return: masks: masks as a `(N, 1, H, W)` array - """ - if rng is None: - rng = np.random - - if self.prop_by_area: - # Choose the proportion of each mask that should be above the threshold - mask_props = rng.uniform(self.prop_range[0], self.prop_range[1], size=(n_masks, self.n_boxes)) - - # Zeros will cause NaNs, so detect and suppres them - zero_mask = mask_props == 0.0 - - if self.random_aspect_ratio: - y_props = np.exp(rng.uniform(low=0.0, high=1.0, size=(n_masks, self.n_boxes)) * np.log(mask_props)) - x_props = mask_props / y_props - else: - y_props = x_props = np.sqrt(mask_props) - fac = np.sqrt(1.0 / self.n_boxes) - y_props *= fac - x_props *= fac - - y_props[zero_mask] = 0 - x_props[zero_mask] = 0 - else: - if self.random_aspect_ratio: - y_props = rng.uniform(self.prop_range[0], self.prop_range[1], size=(n_masks, self.n_boxes)) - x_props = rng.uniform(self.prop_range[0], self.prop_range[1], size=(n_masks, self.n_boxes)) - else: - x_props = y_props = rng.uniform(self.prop_range[0], self.prop_range[1], size=(n_masks, self.n_boxes)) - fac = np.sqrt(1.0 / self.n_boxes) - y_props *= fac - x_props *= fac - - sizes = np.round(np.stack([y_props, x_props], axis=2) * np.array(mask_shape)[None, None, :]) - - if self.within_bounds: - positions = np.round((np.array(mask_shape) - sizes) * rng.uniform(low=0.0, high=1.0, size=sizes.shape)) - rectangles = np.append(positions, positions + sizes, axis=2) - else: - centres = np.round(np.array(mask_shape) * rng.uniform(low=0.0, high=1.0, size=sizes.shape)) - rectangles = np.append(centres - sizes * 0.5, centres + sizes * 0.5, axis=2) - - if self.invert: - masks = np.zeros((n_masks, 1) + mask_shape) - else: - masks = np.ones((n_masks, 1) + mask_shape) - for i, sample_rectangles in enumerate(rectangles): - for y0, x0, y1, x1 in sample_rectangles: - masks[i, 0, int(y0):int(y1), int(x0):int(x1)] = 1 - masks[i, 0, int(y0):int(y1), int(x0):int(x1)] - return masks diff --git a/mpa/modules/experimental/models/segmentors/cutmix_seg_gt_naive.py b/mpa/modules/experimental/models/segmentors/cutmix_seg_gt_naive.py deleted file mode 100644 index aba0c619..00000000 --- a/mpa/modules/experimental/models/segmentors/cutmix_seg_gt_naive.py +++ /dev/null @@ -1,129 +0,0 @@ -import torch -# import torch.nn.functional as F - -from mmseg.models import SEGMENTORS, build_segmentor -from mmseg.models.segmentors.base import BaseSegmentor -from mmseg.ops import resize -import numpy as np -from .cutmix.gt_utils import generate_cutmix_mask -import functools -from collections import OrderedDict -from mpa.utils import logger - - -@SEGMENTORS.register_module() -class CutmixSegGTNaive(BaseSegmentor): - def __init__(self, ori_type=None, unsup_weight=0.1, **kwargs): - print('CutmixSegGTNaive init!') - super(CutmixSegGTNaive, self).__init__() - - cfg = kwargs.copy() - if ori_type == 'SemiSLSegmentor': - cfg['type'] = 'SemiSLSegmentor' - self.align_corners = cfg['decode_head'][-1].align_corners - else: - cfg['type'] = 'EncoderDecoder' - self.align_corners = cfg['decode_head'].align_corners - self.model_s = build_segmentor(cfg) - self.model_t = build_segmentor(cfg) - - self.unsup_weight = unsup_weight - # gt config - self.area_thresh = 0.0001 - self.area_thresh2 = 0.0001 - self.no_pad = False - self.no_slim = False - self.num_classes = kwargs['decode_head'][0]['num_classes'] - self.class_criterion_seg = np.zeros([3, self.num_classes]).astype(float) - - # Hooks for super_type transparent weight load/save - self._register_state_dict_hook(self.state_dict_hook) - self._register_load_state_dict_pre_hook( - functools.partial(self.load_state_dict_pre_hook, self) - ) - - def extract_feat(self, imgs): - return self.model_t.extract_feat(imgs) - - def simple_test(self, img, img_metas, **kwargs): - return self.model_t.simple_test(img, img_metas, **kwargs) - - def aug_test(self, imgs, img_metas, **kwargs): - return self.model_t.aug_test(imgs, img_metas, **kwargs) - - def forward_dummy(self, img, **kwargs): - return self.model_t.forward_dummy(img, **kwargs) - - def forward_train(self, img, img_metas, gt_semantic_seg, **kwargs): - ul_img = kwargs['ul_img'] - ul_img_metas = kwargs['ul_img_metas'] - if 'to_be_inserted_img' in kwargs.keys(): - img_to_be_inserted = kwargs['to_be_inserted_img'] - l_gt = kwargs['to_be_inserted_gt_semantic_seg'] - else: - img_to_be_inserted = img - l_gt = gt_semantic_seg - - if 'sample_cat' in kwargs.keys(): - sample_cat = kwargs['sample_cat'] - else: - sample_cat = [] - for i in range(img.shape[0]): - each_sample_cat = np.random.randint(1, self.num_classes+1) - sample_cat.append(each_sample_cat) - sample_cat = torch.Tensor(sample_cat) - - with torch.no_grad(): - teacher_feat = self.model_t.extract_feat(ul_img) - teacher_logit = self.model_t._decode_head_forward_test(teacher_feat, ul_img_metas) - teacher_logit = resize(input=teacher_logit, - size=ul_img.shape[2:], - mode='bilinear', - align_corners=self.align_corners) - conf_from_teacher, pl_from_teacher = torch.max(torch.softmax(teacher_logit, axis=1), axis=1, keepdim=True) - - masks = generate_cutmix_mask(l_gt.cpu().numpy(), sample_cat.cpu().numpy(), - self.area_thresh, - no_pad=self.no_pad, no_slim=self.no_slim) - if ul_img.is_cuda: - masks = masks.cuda() - ul_img_cutmix = ul_img * (1 - masks) + img_to_be_inserted * masks - # labels for cutmix - tmp = masks * l_gt - pl_from_teacher_cutmixed = (1 - masks) * pl_from_teacher + tmp - pl_from_teacher_cutmixed = pl_from_teacher_cutmixed.long() - - losses = dict() - - x = self.model_s.extract_feat(img) - x_u_cutmixed = self.model_s.extract_feat(ul_img_cutmix) - loss_decode = self.model_s._decode_head_forward_train(x, img_metas, gt_semantic_seg) - loss_decode_u = self.model_s._decode_head_forward_train(x_u_cutmixed, ul_img_metas, pl_from_teacher_cutmixed) - - for key in loss_decode_u.keys(): - losses[key] = (loss_decode[key] + loss_decode_u[key]*self.unsup_weight) - - return losses - - @staticmethod - def state_dict_hook(module, state_dict, *args, **kwargs): - """Redirect teacher model as output state_dict (student as auxilliary) - """ - logger.info('----------------- CutmixSegGTNaive.state_dict_hook() called') - output = OrderedDict() - for k, v in state_dict.items(): - if 'model_t.' in k: - k = k.replace('model_t.', '') - output[k] = v - return output - - @staticmethod - def load_state_dict_pre_hook(module, state_dict, *args, **kwargs): - """Redirect input state_dict to teacher model - """ - logger.info('----------------- CutmixSegGTNaive.load_state_dict_pre_hook() called') - for k in list(state_dict.keys()): - v = state_dict.pop(k) - if 'model_s.' not in k: - state_dict['model_s.'+k] = v - state_dict['model_t.'+k] = v diff --git a/mpa/modules/experimental/models/segmentors/cutmix_seg_naive.py b/mpa/modules/experimental/models/segmentors/cutmix_seg_naive.py deleted file mode 100644 index 7f962a59..00000000 --- a/mpa/modules/experimental/models/segmentors/cutmix_seg_naive.py +++ /dev/null @@ -1,115 +0,0 @@ -import torch - -from mmseg.models import SEGMENTORS, build_segmentor -from mmseg.models.segmentors.base import BaseSegmentor -from mmseg.ops import resize -from .cutmix.mask_gen import BoxMaskGenerator -import functools -from collections import OrderedDict -from mpa.utils import logger - - -@SEGMENTORS.register_module() -class CutmixSegNaive(BaseSegmentor): - def __init__(self, ori_type=None, unsup_weight=0.1, **kwargs): - print('CutmixSegNaive init!') - super(CutmixSegNaive, self).__init__() - - cfg = kwargs.copy() - if ori_type == 'SemiSLSegmentor': - cfg['type'] = 'SemiSLSegmentor' - self.align_corners = cfg['decode_head'][-1].align_corners - else: - cfg['type'] = 'EncoderDecoder' - self.align_corners = cfg['decode_head'].align_corners - self.model_s = build_segmentor(cfg) - self.model_t = build_segmentor(cfg) - - self.unsup_weight = unsup_weight - self.mask_generator = BoxMaskGenerator() - - # Hooks for super_type transparent weight load/save - self._register_state_dict_hook(self.state_dict_hook) - self._register_load_state_dict_pre_hook( - functools.partial(self.load_state_dict_pre_hook, self) - ) - - def extract_feat(self, imgs): - return self.model_t.extract_feat(imgs) - - def simple_test(self, img, img_metas, **kwargs): - return self.model_t.simple_test(img, img_metas, **kwargs) - - def aug_test(self, imgs, img_metas, **kwargs): - return self.model_t.aug_test(imgs, img_metas, **kwargs) - - def forward_dummy(self, img, **kwargs): - return self.model_t.forward_dummy(img, **kwargs) - - def forward_train(self, img, img_metas, gt_semantic_seg, **kwargs): - ul1_img = kwargs['ul_img'] - ul1_img_metas = kwargs['ul_img_metas'] - ul2_img = kwargs['ul2_img'] - ul2_img_metas = kwargs['ul2_img_metas'] - - mask_size = ul1_img.shape[2:] - n_masks = ul1_img.shape[0] - masks = torch.Tensor(self.mask_generator.generate_params(n_masks, mask_size)) - if ul1_img.is_cuda: - masks = masks.cuda() - ul_img_cutmix = (1-masks) * ul1_img + masks * ul2_img - - with torch.no_grad(): - ul1_feat = self.model_t.extract_feat(ul1_img) - ul1_logit = self.model_t._decode_head_forward_test(ul1_feat, ul1_img_metas) - ul1_logit = resize(input=ul1_logit, - size=ul1_img.shape[2:], - mode='bilinear', - align_corners=self.align_corners) - ul1_conf, ul1_pl = torch.max(torch.softmax(ul1_logit, axis=1), axis=1, keepdim=True) - - ul2_feat = self.model_t.extract_feat(ul2_img) - ul2_logit = self.model_t._decode_head_forward_test(ul2_feat, ul2_img_metas) - ul2_logit = resize(input=ul2_logit, - size=ul2_img.shape[2:], - mode='bilinear', - align_corners=self.align_corners) - ul2_conf, ul2_pl = torch.max(torch.softmax(ul2_logit, axis=1), axis=1, keepdim=True) - - pl_cutmixed = (1-masks)*ul1_pl + masks*ul2_pl - pl_cutmixed = pl_cutmixed.long() - - losses = dict() - - x = self.model_s.extract_feat(img) - x_u_cutmixed = self.model_s.extract_feat(ul_img_cutmix) - loss_decode = self.model_s._decode_head_forward_train(x, img_metas, gt_semantic_seg) - loss_decode_u = self.model_s._decode_head_forward_train(x_u_cutmixed, ul1_img_metas, pl_cutmixed) - - for key in loss_decode_u.keys(): - losses[key] = (loss_decode[key] + loss_decode_u[key]*self.unsup_weight) - - return losses - - @staticmethod - def state_dict_hook(module, state_dict, *args, **kwargs): - """Redirect teacher model as output state_dict (student as auxilliary) - """ - logger.info('----------------- CutmixSegNaive.state_dict_hook() called') - output = OrderedDict() - for k, v in state_dict.items(): - if 'model_t.' in k: - k = k.replace('model_t.', '') - output[k] = v - return output - - @staticmethod - def load_state_dict_pre_hook(module, state_dict, *args, **kwargs): - """Redirect input state_dict to teacher model - """ - logger.info('----------------- CutmixSegNaive.load_state_dict_pre_hook() called') - for k in list(state_dict.keys()): - v = state_dict.pop(k) - if 'model_s.' not in k: - state_dict['model_s.'+k] = v - state_dict['model_t.'+k] = v diff --git a/mpa/modules/experimental/models/segmentors/keep_output.py b/mpa/modules/experimental/models/segmentors/keep_output.py deleted file mode 100644 index 284f5049..00000000 --- a/mpa/modules/experimental/models/segmentors/keep_output.py +++ /dev/null @@ -1,77 +0,0 @@ -from mmseg.ops import resize -from mmseg.models import SEGMENTORS, build_segmentor -from mmseg.models.segmentors.base import BaseSegmentor - - -@SEGMENTORS.register_module() -class KeepOutputsWrapper(BaseSegmentor): - def __init__(self, orig_type, **kwargs): - super(KeepOutputsWrapper, self).__init__() - - cfg = kwargs.copy() - cfg['type'] = orig_type - self.segmentor = build_segmentor(cfg) - - if orig_type == 'OCRCascadeEncoderDecoder': - cfg['type'] = 'OCRCascadeEncoderDecoder' - self.align_corners = cfg['decode_head'][-1].align_corners - else: - cfg['type'] = 'EncoderDecoder' - self.align_corners = cfg['decode_head'].align_corners - - print('####################### KeepOutputsWrapper') - - def forward_train(self, img, img_metas, gt_semantic_seg, **kwargs): - segmentor = self.segmentor - - x = segmentor.extract_feat(img) - losses = segmentor._decode_head_forward_train(x, img_metas, gt_semantic_seg) - - logits = segmentor._decode_head_forward_test(x, img_metas) - losses['logits'] = resize(input=logits, - size=img.shape[2:], - mode='bilinear', - align_corners=self.align_corners) - losses['gt_semantic_seg'] = gt_semantic_seg.float() - - return losses - - def simple_test(self, img, img_meta, **kwargs): - return self.segmentor.simple_test(img, img_meta, **kwargs) - - def train_step(self, data_batch, optimizer, **kwargs): - """The iteration step during training. - - This method defines an iteration step during training, except for the - back propagation and optimizer updating, which are done in an optimizer - hook. Note that in some complicated cases or models, the whole process - including back propagation and optimizer updating is also defined in - this method, such as GAN. - - Args: - data (dict): The output of dataloader. - optimizer (:obj:`torch.optim.Optimizer` | dict): The optimizer of - runner is passed to ``train_step()``. This argument is unused - and reserved. - - Returns: - dict: It should contain at least 3 keys: ``loss``, ``log_vars``, - ``num_samples``. - ``loss`` is a tensor for back propagation, which can be a - weighted sum of multiple losses. - ``log_vars`` contains all the variables to be sent to the - logger. - ``num_samples`` indicates the batch size (when the model is - DDP, it means the batch size on each GPU), which is used for - averaging the logs. - """ - losses = self(**data_batch) - loss, log_vars = self._parse_losses(losses) - - outputs = dict( - loss=loss, - log_vars=log_vars, - num_samples=len(data_batch['img'].data)) - outputs['logits'] = losses['logits'] - outputs['gt_semantic_seg'] = losses['gt_semantic_seg'].long() - return outputs diff --git a/mpa/modules/experimental/models/segmentors/mean_teacher_naive.py b/mpa/modules/experimental/models/segmentors/mean_teacher_naive.py deleted file mode 100644 index b681724a..00000000 --- a/mpa/modules/experimental/models/segmentors/mean_teacher_naive.py +++ /dev/null @@ -1,93 +0,0 @@ -from mmseg.models import SEGMENTORS, build_segmentor -from mmseg.models.segmentors.base import BaseSegmentor -from mmseg.ops import resize -import torch -import functools -from collections import OrderedDict -from mpa.utils.logger import get_logger -logger = get_logger() - - -@SEGMENTORS.register_module() -class MeanTeacherNaive(BaseSegmentor): - def __init__(self, ori_type=None, unsup_weight=0.1, **kwargs): - print('MeanTeacherNaive Segmentor init!') - super(MeanTeacherNaive, self).__init__() - - cfg = kwargs.copy() - if ori_type == 'SemiSLSegmentor': - cfg['type'] = 'SemiSLSegmentor' - self.align_corners = cfg['decode_head'][-1].align_corners - else: - cfg['type'] = 'EncoderDecoder' - self.align_corners = cfg['decode_head'].align_corners - self.model_s = build_segmentor(cfg) - self.model_t = build_segmentor(cfg) - - self.unsup_weight = unsup_weight - - # Hooks for super_type transparent weight load/save - self._register_state_dict_hook(self.state_dict_hook) - self._register_load_state_dict_pre_hook( - functools.partial(self.load_state_dict_pre_hook, self) - ) - - def extract_feat(self, imgs): - return self.model_t.extract_feat(imgs) - - def simple_test(self, img, img_metas, **kwargs): - return self.model_t.simple_test(img, img_metas, **kwargs) - - def aug_test(self, imgs, img_metas, **kwargs): - return self.model_t.aug_test(imgs, img_metas, **kwargs) - - def forward_dummy(self, img, **kwargs): - return self.model_t.forward_dummy(img, **kwargs) - - def forward_train(self, img, img_metas, gt_semantic_seg, **kwargs): - ul_img = kwargs['ul_img'] - ul_img_metas = kwargs['ul_img_metas'] - - with torch.no_grad(): - teacher_feat = self.model_t.extract_feat(ul_img) - teacher_logit = self.model_t._decode_head_forward_test(teacher_feat, ul_img_metas) - teacher_logit = resize(input=teacher_logit, - size=ul_img.shape[2:], - mode='bilinear', - align_corners=self.align_corners) - conf_from_teacher, pl_from_teacher = torch.max(torch.softmax(teacher_logit, axis=1), axis=1, keepdim=True) - - losses = dict() - - x = self.model_s.extract_feat(img) - x_u = self.model_s.extract_feat(ul_img) - loss_decode = self.model_s._decode_head_forward_train(x, img_metas, gt_semantic_seg) - loss_decode_u = self.model_s._decode_head_forward_train(x_u, ul_img_metas, pl_from_teacher) - - for key in loss_decode_u.keys(): - losses[key] = (loss_decode[key] + loss_decode_u[key]*self.unsup_weight) - - return losses - - @staticmethod - def state_dict_hook(module, state_dict, *args, **kwargs): - """Redirect teacher model as output state_dict (student as auxilliary) - """ - logger.info('----------------- MeanTeacherNaive.state_dict_hook() called') - output = OrderedDict() - for k, v in state_dict.items(): - if 'model_t.' in k: - k = k.replace('model_t.', '') - output[k] = v - return output - - @staticmethod - def load_state_dict_pre_hook(module, state_dict, *args, **kwargs): - """Redirect input state_dict to teacher model - """ - logger.info('----------------- MeanTeacherNaive.load_state_dict_pre_hook() called') - for k in list(state_dict.keys()): - v = state_dict.pop(k) - if 'model_s.' not in k: - state_dict['model_s.'+k] = v - state_dict['model_t.'+k] = v diff --git a/mpa/modules/experimental/models/segmentors/oversampling.py b/mpa/modules/experimental/models/segmentors/oversampling.py deleted file mode 100644 index 3f674545..00000000 --- a/mpa/modules/experimental/models/segmentors/oversampling.py +++ /dev/null @@ -1,29 +0,0 @@ -import torch - -from mmseg.models import SEGMENTORS, build_segmentor -from mmseg.models.segmentors.base import BaseSegmentor - - -@SEGMENTORS.register_module() -class OverSamplingWrapper(BaseSegmentor): - def __init__(self, orig_type, **kwargs): - super(OverSamplingWrapper, self).__init__() - - cfg = kwargs.copy() - cfg['type'] = orig_type - self.segmentor = build_segmentor(cfg) - - def forward_train(self, img, img_metas, gt_semantic_seg, **kwargs): - img = torch.cat((img, kwargs.pop('os_img')), dim=0) - img_metas += kwargs.pop('os_img_metas') - gt_semantic_seg = torch.cat((gt_semantic_seg, kwargs.pop('os_gt_semantic_seg')), dim=0) - - segmentor = self.segmentor - - x = segmentor.extract_feat(img) - losses = segmentor._decode_head_forward_train(x, img_metas, gt_semantic_seg) - - return losses - - def simple_test(self, img, img_meta, **kwargs): - return self.segmentor.simple_test(img, img_meta, **kwargs) diff --git a/mpa/modules/experimental/models/segmentors/semisl_segmentor.py b/mpa/modules/experimental/models/segmentors/semisl_segmentor.py deleted file mode 100644 index 4ddb988c..00000000 --- a/mpa/modules/experimental/models/segmentors/semisl_segmentor.py +++ /dev/null @@ -1,120 +0,0 @@ -from torch import nn - -from mmseg.core import add_prefix -from mmseg.ops import resize -from mmseg.models import builder, SEGMENTORS -from mmseg.models.segmentors.encoder_decoder import EncoderDecoder - - -@SEGMENTORS.register_module() -class SemiSLSegmentor(EncoderDecoder): - """Cascade Encoder Decoder segmentors. - - CascadeEncoderDecoder almost the same as EncoderDecoder, while decoders of - CascadeEncoderDecoder are cascaded. The output of previous decoder_head - will be the input of next decoder_head. - """ - - def __init__(self, - num_stages, - backbone, - decode_head, - neck=None, - auxiliary_head=None, - train_cfg=None, - test_cfg=None, - pretrained=None, - ori_type=None): - self.num_stages = num_stages - super(SemiSLSegmentor, self).__init__( - backbone=backbone, - decode_head=decode_head, - neck=neck, - auxiliary_head=auxiliary_head, - train_cfg=train_cfg, - test_cfg=test_cfg, - pretrained=pretrained) - - def _init_decode_head(self, decode_head): - """Initialize ``decode_head``""" - assert isinstance(decode_head, list) - assert len(decode_head) == self.num_stages - self.decode_head = nn.ModuleList() - for i in range(self.num_stages): - self.decode_head.append(builder.build_head(decode_head[i])) - self.align_corners = self.decode_head[-1].align_corners - self.num_classes = self.decode_head[-1].num_classes - - def init_weights(self, pretrained=None): - """Initialize the weights in backbone and heads. - - Args: - pretrained (str, optional): Path to pre-trained weights. - Defaults to None. - """ - self.backbone.init_weights(pretrained=pretrained) - for i in range(self.num_stages): - self.decode_head[i].init_weights() - if self.with_auxiliary_head: - if isinstance(self.auxiliary_head, nn.ModuleList): - for aux_head in self.auxiliary_head: - aux_head.init_weights() - else: - self.auxiliary_head.init_weights() - - def encode_decode(self, img, img_metas): - """Encode images with backbone and decode into a semantic segmentation - map of the same size as input.""" - x = self.extract_feat(img) - out = self.decode_head[0].forward_test(x, img_metas, self.test_cfg) - for i in range(1, self.num_stages): - out = self.decode_head[i].forward_test(x, out, img_metas, - self.test_cfg) - out = resize( - input=out, - size=img.shape[2:], - mode='bilinear', - align_corners=self.align_corners) - return out - - def _decode_head_forward_test(self, x, img_metas): - """Run forward function and calculate loss for decode head in - inference.""" - out = self.decode_head[0].forward_test(x, img_metas, self.test_cfg) - for i in range(1, self.num_stages): - out = self.decode_head[i].forward_test(x, out, img_metas, - self.test_cfg) - return out - - def _decode_head_forward_train(self, x, img_metas, gt_semantic_seg): - """Run forward function and calculate loss for decode head in - training.""" - - losses = dict() - - loss_decode = self.decode_head[0].forward_train( - x, img_metas, gt_semantic_seg, self.train_cfg) - - losses.update(add_prefix(loss_decode, 'decode_0')) - - for i in range(1, self.num_stages): - # forward test again, maybe unnecessary for most methods. - if isinstance(x, dict): - prev_outputs = self.decode_head[i - 1].forward_test( - x['x'], img_metas, self.test_cfg) - loss_decode = self.decode_head[i].forward_train( - x['x'], prev_outputs, img_metas, gt_semantic_seg['gt'], self.train_cfg) - losses.update(add_prefix(loss_decode, f'decode_x_{i}')) - prev_outputs = self.decode_head[i - 1].forward_test( - x['x_u'], img_metas, self.test_cfg) - loss_decode = self.decode_head[i].forward_train( - x['x_u'], prev_outputs, img_metas, gt_semantic_seg['pseudo_label'], self.train_cfg) - losses.update(add_prefix(loss_decode, f'decode_x_u_{i}')) - else: - prev_outputs = self.decode_head[i - 1].forward_test( - x, img_metas, self.test_cfg) - loss_decode = self.decode_head[i].forward_train( - x, prev_outputs, img_metas, gt_semantic_seg, self.train_cfg) - losses.update(add_prefix(loss_decode, f'decode_{i}')) - - return losses diff --git a/mpa/modules/experimental/utils/pseudo_nms.py b/mpa/modules/experimental/utils/pseudo_nms.py deleted file mode 100644 index 2a02c218..00000000 --- a/mpa/modules/experimental/utils/pseudo_nms.py +++ /dev/null @@ -1,63 +0,0 @@ -import torch -from mmdet.ops.nms import batched_nms - - -def pseudo_multiclass_nms(multi_bboxes, - multi_scores, - score_thr, - nms_cfg, - max_num=-1, - score_factors=None): - """NMS for multi-class bboxes w/ class probability output - - Coped from mmdet/core/post_processing/bbox_nms.py - to augment NMS output w/ class probabilities as pseudo labels - """ - num_classes = multi_scores.size(1) - 1 - # exclude background category - if multi_bboxes.shape[1] > 4: - bboxes = multi_bboxes.view(multi_scores.size(0), -1, 4) - else: - bboxes = multi_bboxes[:, None].expand( - multi_scores.size(0), num_classes, 4) - scores = multi_scores[:, :-1] - - # filter out boxes with low scores - valid_mask = scores > score_thr - - # We use masked_select for ONNX exporting purpose, - # which is equivalent to bboxes = bboxes[valid_mask] - # (TODO): as ONNX does not support repeat now, - # we have to use this ugly code - bboxes = torch.masked_select( - bboxes, - torch.stack((valid_mask, valid_mask, valid_mask, valid_mask), - -1)).view(-1, 4) - if score_factors is not None: - scores = scores * score_factors[:, None] - scores = torch.masked_select(scores, valid_mask) - labels = valid_mask.nonzero(as_tuple=False)[:, 1] - - if bboxes.numel() == 0: - bboxes = multi_bboxes.new_zeros((0, 5)) - labels = multi_bboxes.new_zeros((0, ), dtype=torch.long) - - if torch.onnx.is_in_onnx_export(): - raise RuntimeError('[ONNX Error] Can not record NMS ' - 'as it has not been executed this time') - return bboxes, labels - - dets, keep = batched_nms(bboxes, scores, labels, nms_cfg) - - if max_num > 0: - dets = dets[:max_num] - keep = keep[:max_num] - - # ----------------- Attach corresponding class probs as pseudo labels - valid_prob_indices = valid_mask.nonzero(as_tuple=False)[:, 0] # Prob row indices w/ any valid score - valid_probs = multi_scores[valid_prob_indices] # Gathering valid class probs - class_probs = valid_probs[keep] # Again gathering non-suppressed class probs - dets = torch.cat((dets, class_probs), dim=1) - # ----------------- - - return dets, labels[keep] diff --git a/mpa/modules/omz/__init__.py b/mpa/modules/omz/__init__.py deleted file mode 100644 index 7d2dea45..00000000 --- a/mpa/modules/omz/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# flake8: noqa -from . import models diff --git a/mpa/modules/omz/custom_layers/__init__.py b/mpa/modules/omz/custom_layers/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/mpa/modules/omz/custom_layers/base_layers.py b/mpa/modules/omz/custom_layers/base_layers.py deleted file mode 100644 index aad58104..00000000 --- a/mpa/modules/omz/custom_layers/base_layers.py +++ /dev/null @@ -1,369 +0,0 @@ -# import numpy as np -import torch -# import torchvision.transforms as transforms -# from torch.nn import * -from torch import nn -from torch.autograd import Variable - - -class Input(nn.Module): - def __init__(self, dims): - super(Input, self).__init__() - self.dims = dims - - def __repr__(self): - return 'Input %s' % self.dims - - def forward(self, inputs): - return inputs - - -class Label(nn.Module): - def __init__(self): - super(Label, self).__init__() - - def __repr__(self): - return 'Labels' - - def forward(self, labels): - return labels - - -class Const(nn.Module): - def __init__(self, data): - super(Const, self).__init__() - # self.data = torch.nn.Parameter(data=torch.from_numpy(data), requires_grad=False) - self.data = torch.from_numpy(data) - # self.register_buffer('_data', self.data) - # self.register_buffer('data', torch.from_numpy(data)) - - def __repr__(self): - return 'Const' - - def forward(self): - return self.data - - -class ScaleShift(nn.Module): - def __init__(self, dims, is_trainable=False): - super(ScaleShift, self).__init__() - self.dims = dims - self.weight = torch.nn.Parameter(data=torch.zeros(dims[1], dtype=torch.float32), requires_grad=is_trainable) - self.bias = torch.nn.Parameter(data=torch.zeros(dims[1], dtype=torch.float32), requires_grad=is_trainable) - # self.register_buffer('weight', torch.zeros(dims[1], dtype=torch.float32)) - # self.register_buffer('bias', torch.zeros(dims[1], dtype=torch.float32)) - - def __repr__(self): - return 'ScaleShift %s' % self.dims - - def forward(self, inputs): - return self.weight[None, :, None, None]*inputs + self.bias[None, :, None, None] - - -class Padding(nn.Module): - def __init__(self, padding): - super(Padding, self).__init__() - self.padding = padding # (x1,x2,y1,y2) - - def __repr__(self): - return 'Padding %s' % self.padding - - def forward(self, inputs): - return nn.functional.pad(inputs, self.padding) - - -class AdaptiveBatchNorm2d(nn.BatchNorm2d): - def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True): - super().__init__(num_features, eps, momentum, affine, track_running_stats) - self.initialized = False - self.num_init_iter = 0 - self.max_init_iter = 10 - - def forward(self, input): - output = super().forward(input) - - if self.training and not self.initialized: - # Use linear output at the first few iteration - output = input*self.weight[None, :, None, None] + self.bias[None, :, None, None] - # output = input*self.weight.detach()[None,:,None, None] + self.bias.detach()[None,:,None, None] - self.num_init_iter += 1 - if self.num_init_iter >= self.max_init_iter: - # Adapt weight & bias using the first batch statistics to undo normalization approximately - self.weight.data = self.weight.data * self.running_var - self.bias.data = self.bias.data + (self.running_mean/(self.running_var + self.eps)) - self.initialized = True - - return output - - -class Conv2dPadBN(nn.Module): - def __init__(self, in_channels, out_channels, kernel_size, stride, - padding, dilation, groups, isbias, isBN, training): - super(Conv2dPadBN, self).__init__() - self.isbias = isbias - self.isBN = isBN - self.training = training - self.pad = Padding(padding) - self.conv = nn.Conv2d(in_channels=in_channels, - out_channels=out_channels, - kernel_size=kernel_size, stride=stride, - dilation=dilation, groups=groups, bias=isbias) - self.weight = self.conv.weight - - if isbias: - self.bias = self.conv.bias - if isBN: - if self.training: - self.bn = AdaptiveBatchNorm2d(out_channels, eps=1.0e-10).train() - else: - self.bn = AdaptiveBatchNorm2d(out_channels, eps=1.0e-10).eval() - nn.init.ones_(self.bn.weight.data) - nn.init.zeros_(self.bn.bias.data) - nn.init.ones_(self.bn.running_var) - nn.init.zeros_(self.bn.running_mean) - - def __repr__(self): - return 'Padding + Conv2D + BN' - - def forward(self, inputs): - x = self.pad(inputs) - x = self.conv(x) - if self.isbias and self.isBN: - x = self.bn(x) - return x - - -class Conv2dBN(nn.Module): - def __init__(self, in_channels, out_channels, kernel_size, stride, - padding, dilation, groups, isbias, isBN, training): - super(Conv2dBN, self).__init__() - self.isbias = isbias - self.isBN = isBN - self.training = training - self.conv = nn.Conv2d(in_channels=in_channels, - out_channels=out_channels, - kernel_size=kernel_size, stride=stride, - padding=padding, - dilation=dilation, groups=groups, bias=isbias) - self.weight = self.conv.weight - - if isbias: - self.bias = self.conv.bias - if isBN: - if self.training: - self.bn = AdaptiveBatchNorm2d(out_channels, eps=1.0e-10).train() - else: - self.bn = AdaptiveBatchNorm2d(out_channels, eps=1.0e-10).eval() - nn.init.ones_(self.bn.weight.data) - nn.init.zeros_(self.bn.bias.data) - nn.init.ones_(self.bn.running_var) - nn.init.zeros_(self.bn.running_mean) - - def __repr__(self): - return repr(self.conv) + '\n ' + repr(self.bn) - - def forward(self, inputs): - x = self.conv(inputs) - if self.isbias and self.isBN: - x = self.bn(x) - return x - - -class MaxPool2DPad(nn.Module): - def __init__(self, kernel_size, stride, padding, ceil_mode): - super(MaxPool2DPad, self).__init__() - self.pad = Padding(padding) - self.maxpool = nn.MaxPool2d(kernel_size=kernel_size, stride=stride, ceil_mode=ceil_mode) - - def __repr__(self): - return 'Padding + MaxPooling2D' - - def forward(self, inputs): - x = self.pad(inputs) - return self.maxpool(x) - - -class AvgPool2DPad(nn.Module): - def __init__(self, kernel_size, stride, padding, ceil_mode): - super(AvgPool2DPad, self).__init__() - self.pad = Padding(padding) - self.avgpool = nn.AvgPool2d(kernel_size=kernel_size, stride=stride, ceil_mode=ceil_mode) - - def __repr__(self): - return 'Padding + AvgPooling2D' - - def forward(self, inputs): - x = self.pad(inputs) - return self.avgpool(x) - - -class Eltwise(nn.Module): - def __init__(self, operation='+'): - super(Eltwise, self).__init__() - self.operation = operation - - def __repr__(self): - return 'Eltwise %s' % self.operation - - def forward(self, inputs): - if self.operation == '+' or self.operation == 'sum': - x = inputs[0] - for i in range(1, len(inputs)): - x = x + inputs[i] - elif self.operation == 'prod' or self.operation == 'mul': - x = inputs[0] - for i in range(1, len(inputs)): - x = torch.mul(x, inputs[i]) - elif self.operation == '/' or self.operation == 'div': - x = inputs[0] - for i in range(1, len(inputs)): - x = x / inputs[i] - elif self.operation == 'max': - x = inputs[0] - for i in range(1, len(inputs)): - x = torch.max(x, inputs[i]) - else: - print('forward Eltwise, unknown operator') - return x - - -class Concat(nn.Module): - def __init__(self, axis): - super(Concat, self).__init__() - self.axis = axis - - def __repr__(self): - return 'Concat(axis=%d)' % self.axis - - def forward(self, inputs): - return torch.cat(inputs, self.axis) - - -class Flatten(nn.Module): - def __init__(self, axis): - super(Flatten, self).__init__() - self.axis = axis - - def __repr__(self): - return 'Flatten(axis=%d)' % self.axis - - def forward(self, x): - left_size = 1 - for i in range(self.axis): - left_size = x.size(i) * left_size - return x.view(left_size, -1).contiguous() - - -class Reshape(nn.Module): - def __init__(self, dims, from_const=True, post_nms_topn=0): - super(Reshape, self).__init__() - self.dims = dims - self.from_const = from_const - self.post_nms_topn = int(post_nms_topn) - if self.from_const and self.post_nms_topn != 0 and (self.dims[0] >= self.post_nms_topn): - self.ratio = int(self.dims[0]/self.post_nms_topn) - else: - self.ratio = 1 - - def __repr__(self): - return 'Reshape(dims=%s)' % (self.dims) - - def forward(self, x): - if self.from_const: - x = x[0] - orig_dims = x.size() - new_dims = [orig_dims[i] if self.dims[i] == 0 else self.dims[i] for i in range(len(self.dims))] - if self.from_const and (self.dims[0] >= self.post_nms_topn): - new_dims[0] = orig_dims[0]*self.ratio - # print('.......origin: ', orig_dims) - # print('.......new: ', new_dims) - - return x.reshape(new_dims) - - -class Crop(nn.Module): - def __init__(self, axis, offset, dim): - super(Crop, self).__init__() - self.axis = axis - self.offset = offset - self.ref = dim - - def __repr__(self): - return 'Crop()' - - def forward(self, x): - idx = 0 - ref = x.shape - x = torch.from_numpy(x) - for axis in self.axis: - ref_size = ref[idx] - offset = int(self.offset[idx]) - indices = torch.arange(offset, ref_size).long() - # indices = x.data.new().resize_(indices.size()).copy_(indices).long() - x = torch.index_select(x, int(axis), Variable(indices)) - idx += 1 - return x - - -class Power(nn.Module): - def __init__(self, power, scale, shift): - super(Power, self).__init__() - self.power = power - self.scale = scale - self.shift = shift - - def __repr__(self): - return 'Power(power=%d, scale=%d, shift=%d)' % (self.power, self.scale, self.shift) - - def forward(self, x): - x = torch.pow(x, self.power) - x = x * self.scale - x = x + self.shift - return x - - -class Permute(nn.Module): - def __init__(self, order): - super(Permute, self).__init__() - self.order = order - - def __repr__(self): - return 'Permute' - - def forward(self, x): - return x.permute(self.order) - - -class Interpolate(nn.Module): - def __init__(self, align_corners, pad_beg, pad_end): - super(Interpolate, self).__init__() - if align_corners == 1: - self.align_corners = True - else: - self.align_corners = False - self.pad_beg = pad_beg - self.pad_end = pad_end - - def __repr__(self): - return 'Interpolate' - - def forward(self, x): - size = tuple((int(x.shape[2]*2), int(x.shape[3]*2))) - return nn.functional.interpolate(x, size, mode='bilinear', align_corners=self.align_corners) - - -class GAPooling(nn.Module): - def __init__(self, stride, padding, ceiling): - super(GAPooling, self).__init__() - self.stride = stride - self.padding = padding - self.ceiling = ceiling - - def __repr__(self): - return 'GAPooling' - - def forward(self, x): - kernel_size = (int(x.shape[2]), int(x.shape[3])) - return nn.functional.avg_pool2d(x, kernel_size=kernel_size, stride=self.stride, - padding=self.padding, ceil_mode=self.ceiling) - # return x.mean([2,3]) diff --git a/mpa/modules/omz/models/__init__.py b/mpa/modules/omz/models/__init__.py deleted file mode 100644 index 3523c40e..00000000 --- a/mpa/modules/omz/models/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -# flake8: noqa -from . import anchor -from . import backbones -from . import dense_heads \ No newline at end of file diff --git a/mpa/modules/omz/models/anchor/__init__.py b/mpa/modules/omz/models/anchor/__init__.py deleted file mode 100644 index 34e430cb..00000000 --- a/mpa/modules/omz/models/anchor/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# from .anchor_generator import SSDAnchorGeneratorClustered diff --git a/mpa/modules/omz/models/anchor/anchor_generator.py b/mpa/modules/omz/models/anchor/anchor_generator.py deleted file mode 100644 index 9d603743..00000000 --- a/mpa/modules/omz/models/anchor/anchor_generator.py +++ /dev/null @@ -1,47 +0,0 @@ -import torch -from torch.nn.modules.utils import _pair - -from mmdet.core.anchor.builder import ANCHOR_GENERATORS -from mmdet.core.anchor.anchor_generator import AnchorGenerator - - -@ANCHOR_GENERATORS.register_module() -class SSDAnchorGeneratorClusteredOMZ(AnchorGenerator): - - def __init__(self, - strides, - widths, - heights): - self.strides = [_pair(stride) for stride in strides] - self.widths = widths - self.heights = heights - self.centers = [(stride / 2., stride / 2.) for stride in strides] - - self.center_offset = 0 - self.base_anchors = self.gen_base_anchors() - - def gen_base_anchors(self): - multi_level_base_anchors = [] - for widths, heights, centers in zip(self.widths, self.heights, self.centers): - base_anchors = self.gen_single_level_base_anchors( - ws=torch.Tensor(widths), - hs=torch.Tensor(heights), - center=torch.Tensor(centers)) - multi_level_base_anchors.append(base_anchors) - return multi_level_base_anchors - - def gen_single_level_base_anchors(self, - ws, - hs, - center): - x_center, y_center = center - - # use float anchor and the anchor's center is aligned with the - # pixel center - base_anchors = [ - x_center - 0.5 * ws, y_center - 0.5 * hs, x_center + 0.5 * ws, - y_center + 0.5 * hs - ] - base_anchors = torch.stack(base_anchors, dim=-1) - - return base_anchors diff --git a/mpa/modules/omz/models/backbones/__init__.py b/mpa/modules/omz/models/backbones/__init__.py deleted file mode 100644 index 756232b7..00000000 --- a/mpa/modules/omz/models/backbones/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -# flake8: noqa -from .omz_backbone_det import OmzBackboneDet -from .omz_backbone_cls import OmzBackboneCls diff --git a/mpa/modules/omz/models/backbones/omz_backbone.py b/mpa/modules/omz/models/backbones/omz_backbone.py deleted file mode 100644 index a25565f3..00000000 --- a/mpa/modules/omz/models/backbones/omz_backbone.py +++ /dev/null @@ -1,74 +0,0 @@ -import torch -# from torch import nn - -# import tlt.core.base.model -from mpa.modules.omz.utils.omz_to_torch import OMZModel, generate_moduledict -from mpa.modules.omz.utils.omz_utils import get_net, get_layers_list -# from tlt.core.base.builder import Builder -# from tlt.core.base.model import Backbone - - -# @Builder.register -class OmzBackbone(OMZModel): - def __init__(self, mode, model_path, last_layer_name, normalized_img_input, **kwargs): - moduledict, layers_info, parents_info = \ - self._parse_model_xml(model_path, last_layer_name, mode == 'train', - mode == 'export', normalized_img_input=normalized_img_input) - super().__init__( - moduledict=moduledict, - layers_info=layers_info, - parents_info=parents_info, - **kwargs) - self.feature_dims = [] - - @staticmethod - def _parse_model_xml(model_path, last_layer_name, training, exporting, normalized_img_input): - omz_net = get_net(model_path) - layers = get_layers_list(omz_net.layers, omz_net.inputs, omz_net.outputs, 'all') - module_dict, _, layers_dict, parents_dict = generate_moduledict( - layers, num_classes=0, training=training, is_export=exporting, - last_layer=last_layer_name, normalized_img_input=normalized_img_input) - return module_dict, layers_dict, parents_dict - - def forward(self, *inputs): - # inputs = inputs*255 # OMZ models take 0~255 inputs - feat = super().forward(inputs) - if len(self.feature_dims) == 0: - self.feature_dims = [feature.shape[1] for feature in self.featuredict.values()] - return feat - - def get_feature_dims(self): - if len(self.feature_dims) == 0: - with torch.no_grad(): - self.forward(torch.zeros(1, 3, 64, 64)) - return self.feature_dims - - -# @Builder.register -class OmzBackboneDummy(OmzBackbone): - - def forward(self, *inputs): - image = inputs[0] - dummy_im_info = torch.Tensor([1, 1, 1.0, 1.0]) - OMZModel.forward(self, [image, dummy_im_info]) - features = list(self.featuredict.values()) - if len(self.feature_dims) == 0: - for feature in features: - if len(feature.shape) > 1: - self.feature_dims.append(feature.shape[1]) - return features - - -# @Builder.register -class OmzBackboneFRCNN(OmzBackbone): - - def forward(self, *inputs): - image = inputs[0] - dummy_im_info = torch.Tensor([1, 1, 1.0, 1.0]) - feat = OMZModel.forward(self, [image, dummy_im_info]) - features = list(self.featuredict.values()) - if len(self.feature_dims) == 0: - for feature in features: - if len(feature.shape) > 1: - self.feature_dims.append(feature.shape[1]) - return feat diff --git a/mpa/modules/omz/models/backbones/omz_backbone_cls.py b/mpa/modules/omz/models/backbones/omz_backbone_cls.py deleted file mode 100644 index fe40dea9..00000000 --- a/mpa/modules/omz/models/backbones/omz_backbone_cls.py +++ /dev/null @@ -1,49 +0,0 @@ -from mmcls.models.builder import BACKBONES -from mpa.modules.omz.models.backbones.omz_backbone import OmzBackbone - - -@BACKBONES.register_module() -class OmzBackboneCls(OmzBackbone): - def __init__(self, - mode, - model_path, - last_layer_name, - frozen_stages=-1, - norm_eval=True, - **kwargs): - super().__init__( - mode=mode, - model_path=model_path, - last_layer_name=last_layer_name, - **kwargs) - - self.frozen_stages = frozen_stages - self.norm_eval = norm_eval - - self._freeze_backbone() - - def forward(self, x): - feature = super().forward(x) - return feature[0] - - def init_weights(self, pretrained=None): - pass - - def train(self, mode=True): - """Convert the model into training mode while keep normalization layer - freezed.""" - super(OmzBackboneCls, self).train(mode) - self._freeze_backbone() - if mode and self.norm_eval: - for m in self.modules(): - # trick: eval have effect on BatchNorm only - classname = m.__class__.__name__ - if classname.find('BatchNorm') != -1: - m.eval() - - def _freeze_backbone(self): - if self.frozen_stages > 0: - for m in self.model.values(): - m.eval() - for param in m.parameters(): - param.requires_grad = False diff --git a/mpa/modules/omz/models/backbones/omz_backbone_det.py b/mpa/modules/omz/models/backbones/omz_backbone_det.py deleted file mode 100644 index ee67dc9d..00000000 --- a/mpa/modules/omz/models/backbones/omz_backbone_det.py +++ /dev/null @@ -1,57 +0,0 @@ -# import tlt.core.base.model -# from tlt.core.omz.omz_to_torch import OMZModel -# from tlt.core.omz.utils.omz_util import OMZModel -# from tlt.core.base.builder import Builder -# from tlt.core.base.model import Backbone -# from tlt.core.backbone.omz import OmzBackboneFRCNN -from mpa.modules.omz.models.backbones.omz_backbone import OmzBackboneFRCNN - -# from ..builder import BACKBONES -from mmdet.models.builder import BACKBONES - - -@BACKBONES.register_module() -class OmzBackboneDet(OmzBackboneFRCNN): - def __init__(self, - mode, - model_path, - last_layer_name, - frozen_stages=-1, - norm_eval=True, - **kwargs): - super().__init__( - mode=mode, - model_path=model_path, - last_layer_name=last_layer_name, - **kwargs) - - self.frozen_stages = frozen_stages - self.norm_eval = norm_eval - - self._freeze_backbone() - - def forward(self, x): # should return a tuple - feature = super().forward(x) - return tuple(feature) - - def init_weights(self, pretrained=None): - pass - - def _freeze_backbone(self): - if self.frozen_stages > 0: - for m in self.model.values(): - m.eval() - for param in m.parameters(): - param.requires_grad = False - - def train(self, mode=True): - """Convert the model into training mode while keep normalization layer - freezed.""" - super(OmzBackboneDet, self).train(mode) - self._freeze_backbone() - if mode and self.norm_eval: - for m in self.modules(): - # trick: eval have effect on BatchNorm only - classname = m.__class__.__name__ - if classname.find('BatchNorm') != -1: - m.eval() diff --git a/mpa/modules/omz/models/dense_heads/__init__.py b/mpa/modules/omz/models/dense_heads/__init__.py deleted file mode 100644 index b28ddc15..00000000 --- a/mpa/modules/omz/models/dense_heads/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -# flake8: noqa -from .omz_rpn_head import OmzRPNHead -from .omz_ssd_head import OmzSSDHead diff --git a/mpa/modules/omz/models/dense_heads/omz_rpn_head.py b/mpa/modules/omz/models/dense_heads/omz_rpn_head.py deleted file mode 100644 index f53fe532..00000000 --- a/mpa/modules/omz/models/dense_heads/omz_rpn_head.py +++ /dev/null @@ -1,57 +0,0 @@ -import torch.nn.functional as F -# from mmcv.ops import batched_nms - -from mmdet.models.builder import HEADS -from mmdet.models.dense_heads.rpn_head import RPNHead - -from mpa.modules.omz.utils.omz_utils import get_net -from mpa.modules.omz.utils.omz_to_torch import gen_convolution - - -@HEADS.register_module() -class OmzRPNHead(RPNHead): - """RPN head. - - Args: - in_channels (int): Number of channels in the input feature map. - """ - - def __init__(self, in_channels, model_path, rpn_conv_name, relu_type, rpn_cls_name, rpn_reg_name, **kwargs): - self.model_path = model_path - self.rpn_conv_name = rpn_conv_name - self.relu_type = relu_type - self.rpn_cls_name = rpn_cls_name - self.rpn_reg_name = rpn_reg_name - super(OmzRPNHead, self).__init__(in_channels, **kwargs) - - def _init_layers(self): - """Initialize layers of the head.""" - omz_net = get_net(self.model_path) - self.rpn_conv = gen_convolution(omz_net.layers[self.rpn_conv_name], training=True) - self.rpn_cls = gen_convolution(omz_net.layers[self.rpn_cls_name], training=True) - self.rpn_reg = gen_convolution(omz_net.layers[self.rpn_reg_name], training=True) - - def init_weights(self): - pass - - def forward_single(self, x): - """Forward feature map of a single scale level.""" - x = self.rpn_conv(x) - if self.relu_type == 'relu6': - x = F.relu6(x, inplace=True) - else: - x = F.relu(x, inplace=True) - rpn_cls_score = self.rpn_cls(x) - rpn_bbox_pred = self.rpn_reg(x) - return rpn_cls_score, rpn_bbox_pred - - def train(self, mode=True): - """Convert the model into training mode while keep normalization layer - freezed.""" - super(OmzRPNHead, self).train(mode) - if mode: - for m in self.modules(): - # trick: eval have effect on BatchNorm only - classname = m.__class__.__name__ - if classname.find('BatchNorm') != -1: - m.eval() diff --git a/mpa/modules/omz/models/dense_heads/omz_ssd_head.py b/mpa/modules/omz/models/dense_heads/omz_ssd_head.py deleted file mode 100644 index b5e98c31..00000000 --- a/mpa/modules/omz/models/dense_heads/omz_ssd_head.py +++ /dev/null @@ -1,73 +0,0 @@ -import torch.nn as nn - -from mmdet.models.builder import HEADS -from mmdet.models.dense_heads.ssd_head import SSDHead - -from mpa.modules.omz.utils.omz_utils import get_net -from mpa.modules.omz.utils.omz_to_torch import generate_moduledict - - -@HEADS.register_module() -class OmzSSDHead(SSDHead): - def __init__(self, model_path, reg_conv_name, cls_conv_name, org_input_size=0, **kwargs): - super(OmzSSDHead, self).__init__(**kwargs) - - # TODO: workaround code to fit scales of anchors with omz model - if org_input_size > 0: - input_size = getattr(self.anchor_generator, 'input_size', org_input_size) - if input_size != org_input_size: - self.anchor_generator.base_sizes = [ - bs*org_input_size/input_size for bs in self.anchor_generator.base_sizes - ] - self.anchor_generator.base_anchors = self.anchor_generator.gen_base_anchors() - - reg_convs = [] - cls_convs = [] - - if isinstance(reg_conv_name, str): - reg_conv_name = [reg_conv_name] - if isinstance(cls_conv_name, str): - cls_conv_name = [cls_conv_name] - - assert len(self.anchor_generator.num_base_anchors) == len(self.in_channels) - assert len(reg_conv_name) == len(self.in_channels) - assert len(cls_conv_name) == len(self.in_channels) - - omz_net = get_net(model_path) - num_classes = self.num_classes + 1 - for i in range(len(self.in_channels)): - in_channels = self.in_channels[i] - num_anchors = self.anchor_generator.num_base_anchors[i] - - _reg_conv_name = reg_conv_name[i] - _cls_conv_name = cls_conv_name[i] - - if isinstance(_reg_conv_name, str): - _reg_conv_name = [_reg_conv_name] - if isinstance(_cls_conv_name, str): - _cls_conv_name = [_cls_conv_name] - - reg_layers = {k: omz_net.layers[k] for k in _reg_conv_name} - cls_layers = {k: omz_net.layers[k] for k in _cls_conv_name} - - reg_module_dict, _, _, _ = generate_moduledict(reg_layers, num_classes) - cls_module_dict, _, _, _ = generate_moduledict(cls_layers, num_classes) - - reg_module_dict.pop('gt_label') - cls_module_dict.pop('gt_label') - - if reg_module_dict[_reg_conv_name[-1]].conv.out_channels == num_anchors*4: - reg_convs.append(nn.Sequential(*reg_module_dict.values())) - else: - reg_convs.append(nn.Conv2d(in_channels, num_anchors*4, kernel_size=3, padding=1)) - - if cls_module_dict[_cls_conv_name[-1]].conv.out_channels == num_anchors*num_classes: - cls_convs.append(nn.Sequential(*cls_module_dict.values())) - else: - cls_convs.append(nn.Conv2d(in_channels, num_anchors*num_classes, kernel_size=3, padding=1)) - - self.reg_convs = nn.ModuleList(reg_convs) - self.cls_convs = nn.ModuleList(cls_convs) - - def init_weights(self): - pass diff --git a/mpa/modules/omz/utils/__init__.py b/mpa/modules/omz/utils/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/mpa/modules/omz/utils/omz_to_torch.py b/mpa/modules/omz/utils/omz_to_torch.py deleted file mode 100644 index 939765d6..00000000 --- a/mpa/modules/omz/utils/omz_to_torch.py +++ /dev/null @@ -1,473 +0,0 @@ -import torch -from collections import OrderedDict -from torch import nn - -# try: -# from openvino import inference_engine as ie # noqa: F401 -# from openvino.inference_engine import IENetwork, IEPlugin, IECore # noqa: F401 -# except Exception as e: -# exception_type = type(e).__name__ -# print("The following error happened while importing Python API module:\n[ {} ] {}".format(exception_type, e)) -# # sys.exit(1) - -from mpa.modules.omz.custom_layers.base_layers import (Conv2dBN, Conv2dPadBN, AdaptiveBatchNorm2d, GAPooling, - MaxPool2DPad, AvgPool2DPad, Input, Label, Const, ScaleShift, - Eltwise, Concat, Reshape, Flatten, Crop, Power, Permute, - Interpolate) -from mpa.modules.omz.utils.omz_utils import get_last_backbone_layer - - -def gen_convolution(layer, training): - params = layer.params - in_channels = layer.in_data[0].shape[1] - out_channels = int(params['output']) - kernel_size = tuple(map(int, params['kernel'].split(','))) - stride = tuple(map(int, params['strides'].split(','))) - pads_begin = list(map(int, params['pads_begin'].split(','))) - pads_end = list(map(int, params['pads_end'].split(','))) - dilation = tuple(map(int, params['dilations'].split(','))) - groups = int(params['group']) - biases_val = layer.blobs.get('biases', None) - auto_pad = params.get('auto_pad', None) - - if biases_val is None: - bias = False - else: - bias = True - - if (layer.in_data[0].shape[2] == 1 and layer.in_data[0].shape[3] == 1): - isBN = False - else: - isBN = True - - if (pads_begin[0] == pads_end[0]) and (pads_begin[1] == pads_end[1]): - if auto_pad == 'valid': - padding = tuple((0, 0)) - elif type(kernel_size) is not tuple and (auto_pad == 'same_upper' or auto_pad == 'same_lower'): - padding = int((kernel_size-1)*0.5) - else: - padding = tuple((pads_begin[0], pads_begin[1])) - convbn = Conv2dBN(in_channels=in_channels, - out_channels=out_channels, - kernel_size=kernel_size, stride=stride, - padding=padding, - dilation=dilation, groups=groups, - isbias=bias, isBN=isBN, training=training) - weights = layer.blobs.get('weights', None) - if weights is not None: - weights = torch.reshape(torch.from_numpy(weights), convbn.weight.shape) - convbn.weight.data = weights - biases = layer.blobs.get('biases', None) - if biases is not None: - biases = torch.reshape(torch.from_numpy(biases), convbn.bias.shape) - convbn.bias.data = biases - return convbn - else: - if auto_pad == 'valid': - padding = tuple((0, 0, 0, 0)) - else: - padding = tuple((pads_begin[1], pads_end[1], pads_begin[0], pads_end[0])) - convpadbn = Conv2dPadBN(in_channels, out_channels, kernel_size, stride, - padding, dilation, groups, - isbias=bias, isBN=isBN, training=training) - weights = layer.blobs.get('weights', None) - if weights is not None: - weights = torch.reshape(torch.from_numpy(weights), convpadbn.weight.shape) - convpadbn.weight.data = weights - biases = layer.blobs.get('biases', None) - if biases is not None: - biases = torch.reshape(torch.from_numpy(biases), convpadbn.bias.shape) - convpadbn.bias.data = biases - return convpadbn - - -def gen_batchnorm(layer, pl, training): - num_features = layer.out_data[0].shape[1] - if training: - bn = AdaptiveBatchNorm2d(num_features=num_features, eps=1.0e-10).train() - else: - bn = AdaptiveBatchNorm2d(num_features=num_features, eps=1.0e-10).eval() - weights = layer.blobs.get('weights', None) - if weights is not None: - weights = torch.reshape(torch.from_numpy(weights), bn.running_var.shape) - bn.weight.data = weights - biases = layer.blobs.get('biases', None) - if biases is not None: - biases = torch.reshape(torch.from_numpy(biases), bn.running_mean.shape) - bn.bias.data = biases - nn.init.ones_(bn.running_var) - nn.init.zeros_(bn.running_mean) - - return bn - - -def gen_elu(layer): - params = layer.params - alpha = float(params['alpha']) - return nn.ELU(alpha=alpha) - - -def gen_relu(layer): - if 'negative_slope' in layer.params: - return nn.LeakyReLU(float(layer.params['negative_slope'])) - else: - return nn.ReLU() - - -def gen_fc(layer): - in_channel = layer.in_data[0].shape[1] - params = layer.params - in_features = in_channel - out_features = int(params['out-size']) - fc = nn.Linear(in_features=in_features, out_features=out_features) - weights = layer.blobs.get('weights', None) - if weights is not None: - weights = torch.reshape(torch.from_numpy(weights), fc.weight.shape) - fc.weight.data = weights - biases = layer.blobs.get('biases', None) - if biases is not None: - biases = torch.reshape(torch.from_numpy(biases), fc.bias.shape) - fc.bias.data = biases - - return fc - - -def gen_pooling(layer): - params = layer.params - kernel_size = tuple(map(int, params['kernel'].split(','))) - stride = tuple(map(int, params['strides'].split(','))) - pads_begin = list(map(int, params['pads_begin'].split(','))) - pads_end = list(map(int, params['pads_end'].split(','))) - ceiling = True - auto_pad = params.get('auto_pad', None) - if params['rounding_type'] == 'floor': - ceiling = False - if (pads_begin[0] == pads_end[0]) and (pads_begin[1] == pads_end[1]): - if auto_pad == 'valid': - padding = tuple((0, 0)) - elif auto_pad == 'same_uppper' or auto_pad == 'same_lower': - padding = int((kernel_size-1)*0.5) - else: - padding = tuple((pads_begin[1], pads_begin[0])) - if params['pool-method'] == 'max': - return nn.MaxPool2d(kernel_size=kernel_size, stride=stride, padding=padding, ceil_mode=ceiling) - elif layer.out_data[0].shape[2] == layer.out_data[0].shape[3] and layer.out_data[0].shape[2] == 1: - return GAPooling(stride, padding, ceiling) - else: - return nn.AvgPool2d(kernel_size=kernel_size, stride=stride, padding=padding, ceil_mode=ceiling) - else: - if auto_pad == 'valid': - padding = tuple((0, 0, 0, 0)) - else: - padding = tuple((pads_begin[1], pads_end[1], pads_begin[0], pads_end[0])) - if params['pool-method'] == 'max': - return MaxPool2DPad(kernel_size=kernel_size, stride=stride, padding=padding, ceil_mode=ceiling) - elif layer.out_data[0].shape[2] == layer.out_data[0].shape[3] and layer.out_data[0].shape[2] == 1: - return GAPooling(stride, ceiling=ceiling) - else: - return AvgPool2DPad(kernel_size=kernel_size, stride=stride, padding=padding, ceil_mode=ceiling) - - -def gen_softmax(layer): - params = layer.params - dim = int(params['axis']) - return nn.Softmax(dim) - - -def gen_input(layer): - dims = layer.out_data[0].shape - return Input(dims) - - -def gen_label(): - return Label() - - -def gen_const(layer): - data = layer.blobs['custom'] - return Const(data) - - -def gen_scaleshift(layer, is_trainable=False): - dims = layer.out_data[0].shape - ss = ScaleShift(dims, is_trainable) - weights = layer.blobs.get('weights', None) - if weights is not None: - weights = torch.reshape(torch.from_numpy(weights), ss.weight.shape) - ss.weight.data = weights - # print('weights: ', weights) - biases = layer.blobs.get('biases', None) - if biases is not None: - biases = torch.reshape(torch.from_numpy(biases), ss.bias.shape) - ss.bias.data = biases - # print('biases: ', biases) - - return ss - - -def gen_eltwise(layer): - params = layer.params - return Eltwise(params['operation']) - - -def gen_concat(layer): - params = layer.params - return Concat(int(params['axis'])) - - -def gen_reshape(layer, parents_l=None, post_nms_topn=0): - if len(layer.parents) > 1: - blobs = parents_l.blobs - dims = blobs['custom'] - return Reshape(dims, post_nms_topn=post_nms_topn) - dims_str = layer.params.get('dim', None) - if dims_str is not None: - dims = list(dims_str.split(',')) - return Reshape(dims, from_const=False) - else: - return Flatten(1) - - -def gen_power(layer): - params = layer.params - power = float(params['power']) - scale = float(params['scale']) - shift = float(params['shift']) - return Power(power, scale, shift) - - -def gen_crop(layer): - params = layer.params - axis = list(params['axis'].split(',')) - offset = list(params['offset'].split(',')) - dim = list(params['dim'].split(',')) - return Crop(axis, offset, dim) - - -def gen_sigmoid(layer): - return nn.Sigmoid() - - -def gen_permute(layer): - params = layer.params - order = list(map(int, params['order'].split(','))) - return Permute(order) - - -def gen_interpolate(layer, parents_l, is_export=False): - params = layer.params - align_corners = int(params['align_corners']) if not is_export else 0 - pad_beg = int(params['pad_beg']) - pad_end = int(params['pad_end']) - return Interpolate(align_corners, pad_beg, pad_end) - - -def gen_mvn(layer): - in_channel = layer.in_data[0].shape[1] - params = layer.params - # across_channels = int(params['across_channels']) - # normalize_variance = int(params['normalize_variance']) - eps = float(params['eps']) - # TODO: consider other parameters - return nn.InstanceNorm2d(num_features=in_channel, eps=eps, affine=True) - - -def generate_moduledict(layers, num_classes, batch_size=32, training=True, is_export=False, - backbone_only=False, stop_layer=None, last_layer=None, normalized_img_input=False): - moduledict = {} - moduledict['gt_label'] = gen_label() - layers_dict = OrderedDict() - parents_dict = {} - post_nms_topn = 0 - normalized_proposal = 0 - # feat_stride = 1 - if backbone_only: - last_layer = get_last_backbone_layer(layers) - bypass_img_input_norm = None - - is_stop = False - for layer in layers.values(): - # print(l.name, l.type, l.params, l.parents, l.blobs.keys()) - module_name = layer.name.replace('.', '_') # ModuleDict does not allow '.' in module name string - if is_stop or (stop_layer is not None and layer.name == stop_layer): - break - if last_layer is not None and layer.name == last_layer: - is_stop = True - if layer.type == 'Input': - moduledict[module_name] = gen_input(layer) - elif layer.type == 'Const': - moduledict[module_name] = gen_const(layer) - elif layer.type == 'ScaleShift': - pl = layers[layer.parents[0]] - if pl.type in ['FullyConnected', 'ScaleShift']: - if batch_size == 1 or not training: - moduledict[module_name] = gen_batchnorm(layer, pl, False) - else: - moduledict[module_name] = gen_batchnorm(layer, pl, True) - elif normalized_img_input and len(layer.parents) == 1 and layers_dict[layer.parents[0]] == 'Image_Input': - if bypass_img_input_norm is not None: - print("bypass_img_input_norm is not None!!!!!!!!!!!!!!!!!!") - bypass_img_input_norm = [layer.children, module_name, layer.parents[0].replace('.', '_')] - continue - else: - moduledict[module_name] = gen_scaleshift(layer) - elif layer.type == 'Convolution': - if (batch_size == 1 or not training): - moduledict[module_name] = gen_convolution(layer, False) - else: - moduledict[module_name] = gen_convolution(layer, True) - elif layer.type == 'FullyConnected': - moduledict[module_name] = gen_fc(layer) - elif layer.type == 'Pooling': - moduledict[module_name] = gen_pooling(layer) - elif layer.type == 'elu': - moduledict[module_name] = gen_elu(layer) - elif layer.type == 'ReLU': - moduledict[module_name] = gen_relu(layer) - elif layer.type == 'Eltwise': - moduledict[module_name] = gen_eltwise(layer) - elif layer.type == 'Concat': - moduledict[module_name] = gen_concat(layer) - elif layer.type == 'Power': - moduledict[module_name] = gen_power(layer) - elif layer.type == 'Crop': - moduledict[module_name] = gen_crop(layer) - elif layer.type == 'Sigmoid': - moduledict[module_name] = gen_sigmoid(layer) - elif layer.type == 'Permute': - moduledict[module_name] = gen_permute(layer) - elif layer.type == 'Interp': - moduledict[module_name] = gen_interpolate(layer, layers[layer.parents[0]], is_export) - elif layer.type == 'Reshape': - if len(layer.parents) > 1: - parents_l = layers[layer.parents[1]] - moduledict[module_name] = gen_reshape(layer, parents_l, post_nms_topn) - else: - moduledict[module_name] = gen_reshape(layer) - elif layer.type == 'SoftMax': - moduledict[module_name] = gen_softmax(layer) - elif layer.type == 'Proposal': - post_nms_topn = layer.params.get('post_nms_topn', 0) - # feat_stride = int(layer.params.get('feat_stride', 1)) - normalized_proposal = int(layer.params.get('normalize', 0)) - continue - elif layer.type == 'MVN': - moduledict[module_name] = gen_mvn(layer) - else: - print('======> extra layers: ', module_name, layer.type, layer.params) - continue - # make network architecture information - if layer.type == 'Input' and len(layer.out_data[0].shape) == 4: - layers_dict[module_name] = 'Image_Input' - else: - layers_dict[module_name] = layer.type - parents_dict[module_name] = list(map(lambda x: x.replace('.', '_'), list(layer.parents))) - - # remove image input normalization layer - if normalized_img_input: - for layer in bypass_img_input_norm[0]: - if layer in parents_dict: - idx = parents_dict[layer].index(bypass_img_input_norm[1]) - parents_dict[layer].insert(idx, bypass_img_input_norm[2]) - parents_dict[layer].remove(bypass_img_input_norm[1]) - - return moduledict, normalized_proposal, layers_dict, parents_dict - - -class OMZModel(nn.Module): - def __init__(self, moduledict, layers_info, parents_info, out_feature_names=[], **kwargs): - super(OMZModel, self).__init__(**kwargs) - self.model = nn.ModuleDict(list(moduledict.items())) - self.layers_info = layers_info - self.parents_info = parents_info - self.featuredict = OrderedDict() - self.softmaxdict = {} - self.logitdict = {} - self.out_feature_names = out_feature_names - - def forward(self, inputs, gt_label=None): - self.featuredict.clear() - self.softmaxdict.clear() - if gt_label is not None: - hidden_layer = self.model['gt_label'] - self.featuredict['gt_label'] = hidden_layer(gt_label) - - input_idx = 1 - for l_name in self.layers_info: - hidden_layer = self.model[l_name] - l_type = self.layers_info[l_name] - parents = list(self.parents_info[l_name]) - if l_type == 'Proposal': - parents.append('gt_label') - input_size = len(parents) - # print('Layer: ', l_name, ': ', l_type, ' <<< ', parents) - if input_size == 1: - input_feature = self.featuredict.get(parents[0], None) - if l_type == 'FullyConnected' and len(input_feature.shape) > 2: - new_dim = [input_feature.shape[0], -1] - input_feature = input_feature.view(new_dim) - feature = hidden_layer(input_feature) - elif input_size > 1: - input_list = [] - for input_name in parents: - input_feature = self.featuredict.get(input_name, None) - if input_feature is None: - print('+++> Missing input feature: ', input_name) - else: - input_list.append(input_feature) - if len(input_list) == 1: - feature = hidden_layer(input_list[0]) - else: - if l_type == 'Proposal': # for PVD structure - input_list[0] = torch.sigmoid(input_list[0]) - feature = hidden_layer(input_list) - elif l_type == 'Image_Input': - feature = hidden_layer(inputs[0]) - elif l_type == 'Input': - feature = hidden_layer(inputs[input_idx]) - input_idx += 1 - else: - feature = hidden_layer() - - # handling output features - if l_type == 'Proposal': - self.featuredict['__RPN_cls_prob'] = self.featuredict[parents[0]] - self.featuredict['__RPN_bbox_pred'] = self.featuredict[parents[1]] - self.featuredict['__FRCNN_loss_input'] = feature - self.featuredict['__Proposal_output'] = feature[0] - feature = feature[0] - elif l_type == 'ROIPooling': - self.featuredict['__ROIPooling_output'] = feature - elif l_type == 'PSROIPooling': - self.featuredict['__PSROIPooling_output'] = feature - elif l_type == 'SoftMax': - self.logitdict[l_name] = self.featuredict[parents[0]] - self.softmaxdict[l_name] = feature - self.featuredict[l_name] = feature - - out_features = [] - if self.out_feature_names: - for l_name in self.out_feature_names: - out_features.append(self.featuredict[l_name]) - else: - out_features.append(feature) - - return out_features - - def get_feature(self, layer_name): - return self.featuredict.get(layer_name, None) - - def get_feature_keys(self): - return self.featuredict.keys() - - def get_softmax_out(self, layer_name): - return self.softmaxdict.get(layer_name, None) - - def get_softmax_out_keys(self): - return self.softmaxdict.keys() - - def get_logit(self, layer_name): - return self.logitdict.get(layer_name, None) - - def get_logit_keys(self): - return self.logitdict.keys() diff --git a/mpa/modules/omz/utils/omz_utils.py b/mpa/modules/omz/utils/omz_utils.py deleted file mode 100644 index b7accbf1..00000000 --- a/mpa/modules/omz/utils/omz_utils.py +++ /dev/null @@ -1,90 +0,0 @@ -import os -import sys - -from collections import OrderedDict - - -def get_net(model: str): - try: - from openvino import inference_engine as ie # noqa: F401 - from openvino.inference_engine import IECore # noqa: F401 - except Exception as e: - exception_type = type(e).__name__ - print("The following error happened while importing Python API module:\n[ {} ] {}".format(exception_type, e)) - sys.exit(1) - - model_xml = model - model_bin = os.path.splitext(model_xml)[0] + ".bin" - iec = IECore() - net = iec.read_network(model=model_xml, weights=model_bin) - return net - - -def get_layers_list(all_layers: dict, inputs: dict, outputs: list, layers: str): - if layers is not None and layers != 'None': - if layers == 'all': - return all_layers - else: - user_layers = [layer.strip() for layer in layers.split(',')] - layers_to_check = [] - for user_layer in user_layers: - if user_layer not in all_layers: - raise Exception("Layer {} doesn't exist in the model".format(user_layer)) - if user_layer in inputs: - raise Exception("Layer {} is input layer. Can not proceed".format(user_layer)) - layers_to_check.append(user_layer) - return layers_to_check - else: - return outputs - - -def get_graph_info_from_openvino(openvino_layers): - layers_dict = OrderedDict() - parents_dict = {} - for layer in openvino_layers.values(): - if layer.type == 'Input' and len(layer.out_data[0].shape) == 4: - layers_dict[layer.name] = 'Image_Input' - else: - layers_dict[layer.name] = layer.type - parents_dict[layer.name] = list(layer.parents) - - return layers_dict, parents_dict - - -def find_root(openvino_layers, leaf_layer, root_layer_name=None, root_id=-1): - parent_layers = leaf_layer.parents - print('***', root_layer_name, ' : ', root_id) - for pl_name in parent_layers: - pl = openvino_layers[pl_name] - if pl.type in ['Input', 'Const']: - continue # im_info - pid = list(openvino_layers.keys()).index(pl_name) - print('...', pl_name, ' : ', pid) - if len(pl.out_data[0].shape) != 4: - return find_root(openvino_layers, pl, root_layer_name, root_id) - elif root_id == -1: - root_id = pid - root_layer_name = pl.name - elif root_id > pid: - leaf_layer = openvino_layers[root_layer_name] - return find_root(openvino_layers, leaf_layer, pl_name, pid) - elif root_id < pid: - return find_root(openvino_layers, pl, root_layer_name, root_id) - else: - break - - return root_layer_name, root_id - - -def get_last_backbone_layer(openvino_layers): - for layer in openvino_layers.values(): - if layer.type == 'Softmax': # classification case - last_backbone_layer, _ = find_root(openvino_layers, layer, None, -1) - break - elif layer.type == 'Proposal': # faster-rcnn case - last_backbone_layer, _ = find_root(openvino_layers, layer, None, -1) - break - else: # recognition case - last_backbone_layer = layer.name - - return last_backbone_layer diff --git a/mpa/modules/omz/utils/utils.py b/mpa/modules/omz/utils/utils.py deleted file mode 100644 index f57cdba5..00000000 --- a/mpa/modules/omz/utils/utils.py +++ /dev/null @@ -1,83 +0,0 @@ -import os - -import numpy as np -import torch -from torch.autograd import Variable - -# from tlt.core.omz.data.voc_data import COLORS, VOC_CLASSES -# from detection.boxes_utils import clip_boxes, bbox_transform_inv, py_cpu_nms -# from detection.boxes_utils import * - - -def to_var(x, *args, **kwargs): - - if isinstance(x, np.ndarray): - x = torch.from_numpy(x) - - if torch.cuda.is_available(): - x = Variable(x, *args, **kwargs).cuda() - else: - x = Variable(x, *args, **kwargs).cpu() - - return x - - -def to_tensor(x, *args, **kwargs): - - if torch.cuda.is_available(): - x = torch.from_numpy(x).cuda() - else: - x = torch.from_numpy(x) - return x - - -def make_name_string(hyparam_dict): - str_result = "" - for i in hyparam_dict.keys(): - str_result = str_result + str(i) + "=" + str(hyparam_dict[i]) + "_" - return str_result[:-1] - - -def read_pickle(path, model, solver): - - try: - files = os.listdir(path) - file_list = [int(file.split('_')[-1].split('.')[0]) for file in files] - file_list.sort() - recent_iter = str(file_list[-1]) - print(recent_iter, path) - - with open(path + "/model_" + recent_iter + ".pkl", "rb") as f: - if torch.cuda.is_available(): - model.load_state_dict(torch.load(f)) - else: - model.load_state_dict(torch.load(f, map_location=lambda storage, loc: storage)) - - with open(path + "/solver_" + recent_iter + ".pkl", "rb") as f: - if torch.cuda.is_available(): - solver.load_state_dict(torch.load(f)) - else: - solver.load_state_dict(torch.load(f, map_location=lambda storage, loc: storage)) - - except Exception as e: - - print("fail try read_pickle", e) - - -def save_pickle(path, epoch, model, solver): - - if not os.path.exists(path): - os.makedirs(path) - - with open(path + "/model_" + str(epoch) + ".pkl", "wb") as f: - torch.save(model.state_dict(), f) - with open(path + "/solver_" + str(epoch) + ".pkl", "wb") as f: - torch.save(solver.state_dict(), f) - - -def save_image(path, iteration, img): - - if not os.path.exists(path): - os.makedirs(path) - - img.save(path + "/img_" + str(iteration) + ".png")