Spaces:
Runtime error
Runtime error
File size: 2,588 Bytes
f549064 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 |
# Copyright (c) OpenMMLab. All rights reserved.
from typing import Callable, Union
import numpy as np
import torch
from mmcls.registry import BATCH_AUGMENTS
class RandomBatchAugment:
"""Randomly choose one batch augmentation to apply.
Args:
augments (Callable | dict | list): configs of batch
augmentations.
probs (float | List[float] | None): The probabilities of each batch
augmentations. If None, choose evenly. Defaults to None.
Example:
>>> import torch
>>> import torch.nn.functional as F
>>> from mmcls.models import RandomBatchAugment
>>> augments_cfg = [
... dict(type='CutMix', alpha=1.),
... dict(type='Mixup', alpha=1.)
... ]
>>> batch_augment = RandomBatchAugment(augments_cfg, probs=[0.5, 0.3])
>>> imgs = torch.rand(16, 3, 32, 32)
>>> label = F.one_hot(torch.randint(0, 10, (16, )), num_classes=10)
>>> imgs, label = batch_augment(imgs, label)
.. note ::
To decide which batch augmentation will be used, it picks one of
``augments`` based on the probabilities. In the example above, the
probability to use CutMix is 0.5, to use Mixup is 0.3, and to do
nothing is 0.2.
"""
def __init__(self, augments: Union[Callable, dict, list], probs=None):
if not isinstance(augments, (tuple, list)):
augments = [augments]
self.augments = []
for aug in augments:
if isinstance(aug, dict):
self.augments.append(BATCH_AUGMENTS.build(aug))
else:
self.augments.append(aug)
if isinstance(probs, float):
probs = [probs]
if probs is not None:
assert len(augments) == len(probs), \
'``augments`` and ``probs`` must have same lengths. ' \
f'Got {len(augments)} vs {len(probs)}.'
assert sum(probs) <= 1, \
'The total probability of batch augments exceeds 1.'
self.augments.append(None)
probs.append(1 - sum(probs))
self.probs = probs
def __call__(self, batch_input: torch.Tensor, batch_score: torch.Tensor):
"""Randomly apply batch augmentations to the batch inputs and batch
data samples."""
aug_index = np.random.choice(len(self.augments), p=self.probs)
aug = self.augments[aug_index]
if aug is not None:
return aug(batch_input, batch_score)
else:
return batch_input, batch_score.float()
|