whcao 3a08db9182
[Feature]Add augments to models/utils (#278)
* add mytrain.py for test

* test before layers

* test attr in layers

* test classifier

* delete mytrain.py

* add rand_bbox_minmax rand_bbox and cutmix_bbox_and_lam to BaseCutMixLayer

* add mixup_prob to BatchMixupLayer

* add cutmixup

* add cutmixup to __init__

* test classifier with cutmixup

* delete some comments

* set mixup_prob default to 1.0

* add cutmixup to classifier

* use cutmixup

* use cutmixup

* fix bugs

* test cutmixup

* move mixup and cutmix to augment

* inherit from BaseAugment

* add BaseAugment

* inherit from BaseAugment

* rename identity.py

* add @

* build augment

* register module

* rename to augment.py

* delete cutmixup.py

* do not inherit from BaseAugment

* add augments

* use augments in classifier

* prob default to 1.0

* add comments

* use augments

* use augments

* assert sum of augmentation probabilities should equal to 1

* augmentation probabilities equal to 1

* calculate Identity prob

* replace xxx with self.xxx

* add comments

* sync with augments

* for BC-breaking

* delete useless comments in mixup.py
2021-06-20 09:44:51 +08:00

30 lines
809 B
Python

import torch.nn.functional as F
from .builder import AUGMENT
@AUGMENT.register_module(name='Identity')
class Identity(object):
"""Change gt_label to one_hot encoding and keep img as the same.
Args:
num_classes (int): The number of classes.
prob (float): MixUp probability. It should be in range [0, 1].
Default to 1.0
"""
def __init__(self, num_classes, prob=1.0):
super(Identity, self).__init__()
assert isinstance(num_classes, int)
assert isinstance(prob, float) and 0.0 <= prob <= 1.0
self.num_classes = num_classes
self.prob = prob
def one_hot(self, gt_label):
return F.one_hot(gt_label, num_classes=self.num_classes)
def __call__(self, img, gt_label):
return img, self.one_hot(gt_label)