def __init__(self, cfg, return_loss=True, **kwargs):
super(YOLOv5, self).__init__(**kwargs)
self.cfg = cfg
self.return_loss = return_loss
self.m, self.save = self._create_model()
self.nms = build_nms(**cfg.test.as_dict())
self.num_anchors = self.cfg.num_anchors
self.num_classes = self.cfg.num_classes
self.g = 0.5
self.off = tf.constant(
[
[0, 0],
[1, 0],
[0, 1],
[-1, 0],
[0, -1], # j,k,l,m
# [1, 1], [1, -1], [-1, 1], [-1, -1], # jk,jm,lk,lm
],
tf.float32) * self.g # offsets
assert "IoU" in cfg.bbox_loss.loss, "Only surpport IoU loss (e.g. IoULoss, GIoU, DIoU, CIoU)."
self.bbox_loss_func = build_loss(**cfg.bbox_loss.as_dict())
self.label_loss_func = build_loss(**cfg.label_loss.as_dict())
self.conf_loss_func = build_loss(**cfg.conf_loss.as_dict())
def __init__(self, cfg, num_classes, **kwargs):
super(FCOSHead, self).__init__(cfg, num_classes=num_classes, **kwargs)
self._make_shared_convs()
self.classifier = tf.keras.layers.Conv2D(
filters=cfg.num_classes,
kernel_size=(3, 3),
strides=(1, 1),
padding="same",
use_bias=True,
bias_initializer=tf.keras.initializers.Constant(-math.log(
(1. - cfg.prior) / cfg.prior)),
name="predicted_class")
self.regressor = tf.keras.layers.Conv2D(filters=4,
kernel_size=(3, 3),
strides=(1, 1),
padding="same",
use_bias=True,
name="predicted_box")
self.centerness = tf.keras.layers.Conv2D(
filters=1,
kernel_size=(3, 3),
strides=(1, 1),
use_bias=True,
kernel_initializer=tf.keras.initializers.RandomNormal(stddev=0.01),
padding="same",
name="predicted_centerness")
self.scales = [
Scale(value=1.0, name="%d" % i)
for i, _ in enumerate(range(cfg.min_level, cfg.max_level + 1))
]
self.normalize_box = cfg.normalize_box
self.centerness_loss_func = build_loss(**cfg.centerness_loss.as_dict())
def __init__(self, cfg, anchor_cfg, num_classes, **kwargs):
super(GFLHead, self).__init__(cfg=cfg,
anchor_cfg=anchor_cfg,
num_classes=num_classes,
**kwargs)
self._make_shared_convs()
self._init_anchor_generators()
self.classifier = tf.keras.layers.Conv2D(
filters=cfg.num_classes,
kernel_size=(3, 3),
strides=(1, 1),
padding="same",
use_bias=True,
bias_initializer=tf.keras.initializers.Constant(-math.log(
(1. - cfg.prior) / cfg.prior)),
name="predicted_class")
self.regressor = tf.keras.layers.Conv2D(filters=4 * (cfg.reg_max + 1),
kernel_size=(3, 3),
strides=(1, 1),
padding="same",
use_bias=True,
name="predicted_box")
self.scales = [
Scale(value=1.0, name="scales/%d" % i)
for i, _ in enumerate(range(cfg.min_level, cfg.max_level + 1))
]
self._num_anchors_per_level = dict()
self.project = tf.linspace(0, cfg.reg_max, cfg.reg_max + 1)
self.dfl_loss_func = build_loss(**cfg.dfl_loss.as_dict())
def __init__(self,
cfg,
test_cfg,
anchor_cfg=None,
num_classes=80,
is_training=True,
data_format="channels_last",
**kwargs):
super(BaseHead, self).__init__(**kwargs)
self.num_classes = num_classes
self.cfg = cfg
self.anchor_cfg = anchor_cfg
self.test_cfg = test_cfg
self.is_training = is_training
self.data_format = data_format
if test_cfg and test_cfg.get("nms") is not None:
self.nms = build_nms(**test_cfg.as_dict())
self.use_sigmoid = True
if cfg.get("use_sigmoid") is not None:
self.use_sigmoid = cfg.use_sigmoid
self._label_dims = num_classes if self.use_sigmoid else num_classes + 1
self.bbox_loss_func = build_loss(**cfg.bbox_loss.as_dict()) if cfg.get(
"bbox_loss") is not None else None
self._use_iou_loss = False
if self.bbox_loss_func is not None:
self._use_iou_loss = "IoU" in cfg.bbox_loss.loss
self.label_loss_func = build_loss(**cfg.label_loss.as_dict(
)) if cfg.get("label_loss") is not None else None
self.sampler = build_sampler(**cfg.sampler.as_dict()) if cfg.get(
"sampler") is not None else None
self.assigner = build_assigner(**cfg.assigner.as_dict()) if cfg.get(
"assigner") is not None else None
self.bbox_decoder = build_decoder(**cfg.bbox_decoder.as_dict(
)) if cfg.get("bbox_decoder") is not None else None
self.bbox_encoder = build_encoder(**cfg.bbox_encoder.as_dict(
)) if cfg.get("bbox_encoder") is not None else None
def __init__(self, cfg, **kwargs):
super(OneNetHead, self).__init__(cfg, **kwargs)
assert cfg.assigner.assigner == "MinCostAssigner"
self.feat = tf.keras.layers.Conv2D(
filters=cfg.feat_dims,
kernel_size=(3, 3),
strides=(1, 1),
padding="same",
use_bias=True,
name="feat1")
self.classifier = tf.keras.layers.Conv2D(
filters=self.num_classes,
kernel_size=(3, 3),
strides=(1, 1),
padding="same",
use_bias=True,
bias_initializer=tf.keras.initializers.Constant(
-math.log((1. - cfg.prior) / cfg.prior)),
name="predicted_class")
self.regressor = tf.keras.layers.Conv2D(
filters=4,
kernel_size=(3, 3),
strides=(1, 1),
padding="same",
use_bias=True,
name="predicted_box")
self.act = build_activation(**cfg.activation.as_dict())
self.assigner = build_assigner(**cfg.assigner.as_dict()) if cfg.get("assigner") is not None else None
self.bbox_loss_func = build_loss(**cfg.bbox_loss.as_dict()) if cfg.get("bbox_loss") is not None else None
self._use_iou_loss = False
if self.bbox_loss_func is not None:
self._use_iou_loss = "IoU" in cfg.bbox_loss.loss
self.label_loss_func = build_loss(**cfg.label_loss.as_dict()) if cfg.get("label_loss") is not None else None
def __init__(self, cfg, num_classes, **kwargs):
super(CenterHeatmapHead, self).__init__(cfg=cfg,
num_classes=num_classes,
**kwargs)
self.hm_loss_func = build_loss(**cfg.hm_loss.as_dict())
self.wh_loss_func = build_loss(**cfg.wh_loss.as_dict())
self.reg_loss_func = build_loss(**cfg.reg_loss.as_dict())
bias_value = -2.19 # -math.log((1 - 0.1) / 0.1)
if cfg.num_stacks > 1:
for i in range(cfg.num_stacks):
hm = tf.keras.Sequential(name="hm/%d" % i)
wh = tf.keras.Sequential(name="wh/%d" % i)
reg = tf.keras.Sequential(name="reg/%d" % i)
for j in range(cfg.repeats):
hm.add(
ConvNormActBlock(
filters=cfg.feat_dims,
kernel_size=3,
data_format=self.data_format,
normalization=cfg.normalization.as_dict()
if cfg.normalization else None,
activation=cfg.activation.as_dict(),
name="%d" % j))
wh.add(
ConvNormActBlock(
filters=cfg.feat_dims,
kernel_size=3,
data_format=self.data_format,
normalization=cfg.normalization.as_dict()
if cfg.normalization else None,
activation=cfg.activation.as_dict(),
name="%d" % j))
reg.add(
ConvNormActBlock(
filters=cfg.feat_dims,
kernel_size=3,
data_format=self.data_format,
normalization=cfg.normalization.as_dict()
if cfg.normalization else None,
activation=cfg.activation.as_dict(),
name="%d" % j))
hm.add(
tf.keras.layers.Conv2D(
filters=num_classes,
kernel_size=1,
data_format=self.data_format,
bias_initializer=tf.keras.initializers.Constant(
bias_value),
name="predicted_hm%d" % cfg.repeats))
wh.add(
tf.keras.layers.Conv2D(filters=2,
kernel_size=1,
data_format=self.data_format,
name="predicted_wh%d" %
cfg.repeats))
reg.add(
tf.keras.layers.Conv2D(filters=2,
kernel_size=1,
data_format=self.data_format,
name="predicted_reg%d" %
cfg.repeats))
setattr(self, "wh%d" % i, wh)
setattr(self, "hm%d" % i, hm)
setattr(self, "reg%d" % i, reg)
else:
hm = tf.keras.Sequential(name="hm")
wh = tf.keras.Sequential(name="wh")
reg = tf.keras.Sequential(name="reg")
for j in range(cfg.repeats):
hm.add(
ConvNormActBlock(filters=cfg.feat_dims,
kernel_size=3,
data_format=self.data_format,
normalization=cfg.normalization.as_dict()
if cfg.normalization else None,
activation=cfg.activation.as_dict(),
name="%d" % j))
wh.add(
ConvNormActBlock(filters=cfg.feat_dims,
kernel_size=3,
data_format=self.data_format,
normalization=cfg.normalization.as_dict()
if cfg.normalization else None,
activation=cfg.activation.as_dict(),
name="%d" % j))
reg.add(
ConvNormActBlock(filters=cfg.feat_dims,
kernel_size=3,
data_format=self.data_format,
normalization=cfg.normalization.as_dict()
if cfg.normalization else None,
activation=cfg.activation.as_dict(),
name="%d" % j))
hm.add(
tf.keras.layers.Conv2D(
filters=num_classes,
kernel_size=1,
data_format=self.data_format,
bias_initializer=tf.keras.initializers.Constant(
bias_value),
name="predicted_hm"))
wh.add(
tf.keras.layers.Conv2D(filters=2,
kernel_size=1,
data_format=self.data_format,
name="predicted_wh"))
reg.add(
tf.keras.layers.Conv2D(filters=2,
kernel_size=1,
data_format=self.data_format,
name="predicted_reg"))
setattr(self, "wh", wh)
setattr(self, "hm", hm)
setattr(self, "reg", reg)
def __init__(self, cfg):
use_mixed_precision = cfg.dtype == "float16"
if use_mixed_precision:
policy = tf.keras.mixed_precision.experimental.Policy("mixed_float16")
tf.keras.mixed_precision.experimental.set_policy(policy)
self.train_dataset = build_dataset(**cfg.train.dataset.as_dict())
self.val_dataset = build_dataset(**cfg.val.dataset.as_dict())
self.model = build_model(**cfg.model.as_dict())
optimizer = build_optimizer(**cfg.optimizer.as_dict())
if cfg.lookahead:
optimizer = LookaheadOptimizer(optimizer, cfg.lookahead.steps, cfg.lookahead.alpha)
if use_mixed_precision:
optimizer = tf.keras.mixed_precision.experimental.LossScaleOptimizer(
optimizer=optimizer, loss_scale= "dynamic")
self.loss_fn = build_loss(**cfg.loss.as_dict())
self.optimizer = optimizer
self.use_mixed_precision = use_mixed_precision
self.cfg = cfg
self.total_train_steps = cfg.learning_rate_scheduler.train_steps
self.learning_rate_scheduler = build_learning_rate_scheduler(
**cfg.learning_rate_scheduler.as_dict())
self.global_step = tf.Variable(initial_value=0,
trainable=False,
name="global_step",
dtype=tf.int64)
self.learning_rate = tf.Variable(initial_value=0,
trainable=False,
name="learning_rate",
dtype=tf.float32)
self.checkpoint = tf.train.Checkpoint(optimizer=self.optimizer, model=self.model)
self.manager = tf.train.CheckpointManager(checkpoint=self.checkpoint,
directory=cfg.checkpoint_dir,
max_to_keep=10)
if os.path.exists(cfg.pretrained_weights_path):
self.model.load_weights(cfg.pretrained_weights_path, by_name=True, skip_mismatch=True)
latest_checkpoint = self.manager.latest_checkpoint
if latest_checkpoint is not None:
try:
steps = int(latest_checkpoint.split("-")[-1])
self.global_step.assign(steps)
except:
self.global_step.assign(0)
self.checkpoint.restore(latest_checkpoint)
tf.print(_time_to_string(), "Restored weights from %s." % latest_checkpoint)
else:
self.global_step.assign(0)
self.summary_writer = tf.summary.create_file_writer(logdir=cfg.summary_dir)
self.log_every_n_steps = cfg.log_every_n_steps
self.save_ckpt_steps = cfg.save_ckpt_steps
self.use_jit = tf.config.optimizer.get_jit() is not None
self.training_loss_metrics = {}
self.val_loss_metrics = {}
self.train_acc_metric = tf.keras.metrics.Accuracy()
self.train_auc_metric = tf.keras.metrics.AUC()
self.val_acc_metric = tf.keras.metrics.Accuracy()
self.val_auc_metric = tf.keras.metrics.AUC()
self._add_graph = True