File size: 2,992 Bytes
e85fecb |
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 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 |
"""
Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""
from typing import Any, Dict, List, Optional
import torch
import torch.nn as nn
import torchvision
import torchvision.transforms.v2 as T
from ...core import GLOBAL_CONFIG, register
from ._transforms import EmptyTransform
torchvision.disable_beta_transforms_warning()
@register()
class Compose(T.Compose):
def __init__(self, ops, policy=None) -> None:
transforms = []
if ops is not None:
for op in ops:
if isinstance(op, dict):
name = op.pop("type")
transform = getattr(
GLOBAL_CONFIG[name]["_pymodule"], GLOBAL_CONFIG[name]["_name"]
)(**op)
transforms.append(transform)
op["type"] = name
elif isinstance(op, nn.Module):
transforms.append(op)
else:
raise ValueError("")
else:
transforms = [
EmptyTransform(),
]
super().__init__(transforms=transforms)
if policy is None:
policy = {"name": "default"}
self.policy = policy
self.global_samples = 0
def forward(self, *inputs: Any) -> Any:
return self.get_forward(self.policy["name"])(*inputs)
def get_forward(self, name):
forwards = {
"default": self.default_forward,
"stop_epoch": self.stop_epoch_forward,
"stop_sample": self.stop_sample_forward,
}
return forwards[name]
def default_forward(self, *inputs: Any) -> Any:
sample = inputs if len(inputs) > 1 else inputs[0]
for transform in self.transforms:
sample = transform(sample)
return sample
def stop_epoch_forward(self, *inputs: Any):
sample = inputs if len(inputs) > 1 else inputs[0]
dataset = sample[-1]
cur_epoch = dataset.epoch
policy_ops = self.policy["ops"]
policy_epoch = self.policy["epoch"]
for transform in self.transforms:
if type(transform).__name__ in policy_ops and cur_epoch >= policy_epoch:
pass
else:
sample = transform(sample)
return sample
def stop_sample_forward(self, *inputs: Any):
sample = inputs if len(inputs) > 1 else inputs[0]
dataset = sample[-1]
cur_epoch = dataset.epoch
policy_ops = self.policy["ops"]
policy_sample = self.policy["sample"]
for transform in self.transforms:
if type(transform).__name__ in policy_ops and self.global_samples >= policy_sample:
pass
else:
sample = transform(sample)
self.global_samples += 1
return sample
|