first commit
This commit is contained in:
陈赣
26
rtdetr_paddle/ppdet/engine/__init__.py
Normal file
26
rtdetr_paddle/ppdet/engine/__init__.py
Normal file
@@ -0,0 +1,26 @@
|
||||
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from . import trainer
|
||||
from .trainer import *
|
||||
|
||||
from . import callbacks
|
||||
from .callbacks import *
|
||||
|
||||
from . import env
|
||||
from .env import *
|
||||
|
||||
__all__ = trainer.__all__ \
|
||||
+ callbacks.__all__ \
|
||||
+ env.__all__
|
||||
557
rtdetr_paddle/ppdet/engine/callbacks.py
Normal file
557
rtdetr_paddle/ppdet/engine/callbacks.py
Normal file
@@ -0,0 +1,557 @@
|
||||
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import os
|
||||
import sys
|
||||
import datetime
|
||||
import six
|
||||
import copy
|
||||
import json
|
||||
|
||||
import paddle
|
||||
import paddle.distributed as dist
|
||||
|
||||
from ppdet.utils.checkpoint import save_model
|
||||
from ppdet.metrics import get_infer_results
|
||||
|
||||
from ppdet.utils.logger import setup_logger
|
||||
logger = setup_logger('ppdet.engine')
|
||||
|
||||
__all__ = [
|
||||
'Callback', 'ComposeCallback', 'LogPrinter', 'Checkpointer',
|
||||
'VisualDLWriter', 'SniperProposalsGenerator'
|
||||
]
|
||||
|
||||
|
||||
class Callback(object):
|
||||
def __init__(self, model):
|
||||
self.model = model
|
||||
|
||||
def on_step_begin(self, status):
|
||||
pass
|
||||
|
||||
def on_step_end(self, status):
|
||||
pass
|
||||
|
||||
def on_epoch_begin(self, status):
|
||||
pass
|
||||
|
||||
def on_epoch_end(self, status):
|
||||
pass
|
||||
|
||||
def on_train_begin(self, status):
|
||||
pass
|
||||
|
||||
def on_train_end(self, status):
|
||||
pass
|
||||
|
||||
|
||||
class ComposeCallback(object):
|
||||
def __init__(self, callbacks):
|
||||
callbacks = [c for c in list(callbacks) if c is not None]
|
||||
for c in callbacks:
|
||||
assert isinstance(
|
||||
c, Callback), "callback should be subclass of Callback"
|
||||
self._callbacks = callbacks
|
||||
|
||||
def on_step_begin(self, status):
|
||||
for c in self._callbacks:
|
||||
c.on_step_begin(status)
|
||||
|
||||
def on_step_end(self, status):
|
||||
for c in self._callbacks:
|
||||
c.on_step_end(status)
|
||||
|
||||
def on_epoch_begin(self, status):
|
||||
for c in self._callbacks:
|
||||
c.on_epoch_begin(status)
|
||||
|
||||
def on_epoch_end(self, status):
|
||||
for c in self._callbacks:
|
||||
c.on_epoch_end(status)
|
||||
|
||||
def on_train_begin(self, status):
|
||||
for c in self._callbacks:
|
||||
c.on_train_begin(status)
|
||||
|
||||
def on_train_end(self, status):
|
||||
for c in self._callbacks:
|
||||
c.on_train_end(status)
|
||||
|
||||
|
||||
class LogPrinter(Callback):
|
||||
def __init__(self, model):
|
||||
super(LogPrinter, self).__init__(model)
|
||||
|
||||
def on_step_end(self, status):
|
||||
if dist.get_world_size() < 2 or dist.get_rank() == 0:
|
||||
mode = status['mode']
|
||||
if mode == 'train':
|
||||
epoch_id = status['epoch_id']
|
||||
step_id = status['step_id']
|
||||
steps_per_epoch = status['steps_per_epoch']
|
||||
training_status = status['training_status']
|
||||
batch_time = status['batch_time']
|
||||
data_time = status['data_time']
|
||||
|
||||
epoches = self.model.cfg.epoch
|
||||
batch_size = self.model.cfg['{}Reader'.format(mode.capitalize(
|
||||
))]['batch_size']
|
||||
|
||||
logs = training_status.log()
|
||||
space_fmt = ':' + str(len(str(steps_per_epoch))) + 'd'
|
||||
if step_id % self.model.cfg.log_iter == 0:
|
||||
eta_steps = (epoches - epoch_id) * steps_per_epoch - step_id
|
||||
eta_sec = eta_steps * batch_time.global_avg
|
||||
eta_str = str(datetime.timedelta(seconds=int(eta_sec)))
|
||||
ips = float(batch_size) / batch_time.avg
|
||||
fmt = ' '.join([
|
||||
'Epoch: [{}]',
|
||||
'[{' + space_fmt + '}/{}]',
|
||||
'learning_rate: {lr:.6f}',
|
||||
'{meters}',
|
||||
'eta: {eta}',
|
||||
'batch_cost: {btime}',
|
||||
'data_cost: {dtime}',
|
||||
'ips: {ips:.4f} images/s',
|
||||
])
|
||||
fmt = fmt.format(
|
||||
epoch_id,
|
||||
step_id,
|
||||
steps_per_epoch,
|
||||
lr=status['learning_rate'],
|
||||
meters=logs,
|
||||
eta=eta_str,
|
||||
btime=str(batch_time),
|
||||
dtime=str(data_time),
|
||||
ips=ips)
|
||||
logger.info(fmt)
|
||||
if mode == 'eval':
|
||||
step_id = status['step_id']
|
||||
if step_id % 100 == 0:
|
||||
logger.info("Eval iter: {}".format(step_id))
|
||||
|
||||
def on_epoch_end(self, status):
|
||||
if dist.get_world_size() < 2 or dist.get_rank() == 0:
|
||||
mode = status['mode']
|
||||
if mode == 'eval':
|
||||
sample_num = status['sample_num']
|
||||
cost_time = status['cost_time']
|
||||
logger.info('Total sample number: {}, average FPS: {}'.format(
|
||||
sample_num, sample_num / cost_time))
|
||||
|
||||
|
||||
class Checkpointer(Callback):
|
||||
def __init__(self, model):
|
||||
super(Checkpointer, self).__init__(model)
|
||||
self.best_ap = -1000.
|
||||
self.save_dir = os.path.join(self.model.cfg.save_dir,
|
||||
self.model.cfg.filename)
|
||||
if hasattr(self.model.model, 'student_model'):
|
||||
self.weight = self.model.model.student_model
|
||||
else:
|
||||
self.weight = self.model.model
|
||||
|
||||
def on_epoch_end(self, status):
|
||||
# Checkpointer only performed during training
|
||||
mode = status['mode']
|
||||
epoch_id = status['epoch_id']
|
||||
weight = None
|
||||
save_name = None
|
||||
if dist.get_world_size() < 2 or dist.get_rank() == 0:
|
||||
if mode == 'train':
|
||||
end_epoch = self.model.cfg.epoch
|
||||
if (
|
||||
epoch_id + 1
|
||||
) % self.model.cfg.snapshot_epoch == 0 or epoch_id == end_epoch - 1:
|
||||
save_name = str(
|
||||
epoch_id) if epoch_id != end_epoch - 1 else "model_final"
|
||||
weight = self.weight.state_dict()
|
||||
elif mode == 'eval':
|
||||
if 'save_best_model' in status and status['save_best_model']:
|
||||
for metric in self.model._metrics:
|
||||
map_res = metric.get_results()
|
||||
eval_func = "ap"
|
||||
if 'pose3d' in map_res:
|
||||
key = 'pose3d'
|
||||
eval_func = "mpjpe"
|
||||
elif 'bbox' in map_res:
|
||||
key = 'bbox'
|
||||
elif 'keypoint' in map_res:
|
||||
key = 'keypoint'
|
||||
else:
|
||||
key = 'mask'
|
||||
if key not in map_res:
|
||||
logger.warning("Evaluation results empty, this may be due to " \
|
||||
"training iterations being too few or not " \
|
||||
"loading the correct weights.")
|
||||
return
|
||||
if map_res[key][0] >= self.best_ap:
|
||||
self.best_ap = map_res[key][0]
|
||||
save_name = 'best_model'
|
||||
weight = self.weight.state_dict()
|
||||
logger.info("Best test {} {} is {:0.3f}.".format(
|
||||
key, eval_func, abs(self.best_ap)))
|
||||
if weight:
|
||||
if self.model.use_ema:
|
||||
exchange_save_model = status.get('exchange_save_model',
|
||||
False)
|
||||
if not exchange_save_model:
|
||||
# save model and ema_model
|
||||
save_model(
|
||||
status['weight'],
|
||||
self.model.optimizer,
|
||||
self.save_dir,
|
||||
save_name,
|
||||
epoch_id + 1,
|
||||
ema_model=weight)
|
||||
else:
|
||||
# save model(student model) and ema_model(teacher model)
|
||||
# in DenseTeacher SSOD, the teacher model will be higher,
|
||||
# so exchange when saving pdparams
|
||||
student_model = status['weight'] # model
|
||||
teacher_model = weight # ema_model
|
||||
save_model(
|
||||
teacher_model,
|
||||
self.model.optimizer,
|
||||
self.save_dir,
|
||||
save_name,
|
||||
epoch_id + 1,
|
||||
ema_model=student_model)
|
||||
del teacher_model
|
||||
del student_model
|
||||
else:
|
||||
save_model(weight, self.model.optimizer, self.save_dir,
|
||||
save_name, epoch_id + 1)
|
||||
|
||||
|
||||
class WiferFaceEval(Callback):
|
||||
def __init__(self, model):
|
||||
super(WiferFaceEval, self).__init__(model)
|
||||
|
||||
def on_epoch_begin(self, status):
|
||||
assert self.model.mode == 'eval', \
|
||||
"WiferFaceEval can only be set during evaluation"
|
||||
for metric in self.model._metrics:
|
||||
metric.update(self.model.model)
|
||||
sys.exit()
|
||||
|
||||
|
||||
class VisualDLWriter(Callback):
|
||||
"""
|
||||
Use VisualDL to log data or image
|
||||
"""
|
||||
|
||||
def __init__(self, model):
|
||||
super(VisualDLWriter, self).__init__(model)
|
||||
|
||||
assert six.PY3, "VisualDL requires Python >= 3.5"
|
||||
try:
|
||||
from visualdl import LogWriter
|
||||
except Exception as e:
|
||||
logger.error('visualdl not found, plaese install visualdl. '
|
||||
'for example: `pip install visualdl`.')
|
||||
raise e
|
||||
self.vdl_writer = LogWriter(
|
||||
model.cfg.get('vdl_log_dir', 'vdl_log_dir/scalar'))
|
||||
self.vdl_loss_step = 0
|
||||
self.vdl_mAP_step = 0
|
||||
self.vdl_image_step = 0
|
||||
self.vdl_image_frame = 0
|
||||
|
||||
def on_step_end(self, status):
|
||||
mode = status['mode']
|
||||
if dist.get_world_size() < 2 or dist.get_rank() == 0:
|
||||
if mode == 'train':
|
||||
training_status = status['training_status']
|
||||
for loss_name, loss_value in training_status.get().items():
|
||||
self.vdl_writer.add_scalar(loss_name, loss_value,
|
||||
self.vdl_loss_step)
|
||||
self.vdl_loss_step += 1
|
||||
elif mode == 'test':
|
||||
ori_image = status['original_image']
|
||||
result_image = status['result_image']
|
||||
self.vdl_writer.add_image(
|
||||
"original/frame_{}".format(self.vdl_image_frame), ori_image,
|
||||
self.vdl_image_step)
|
||||
self.vdl_writer.add_image(
|
||||
"result/frame_{}".format(self.vdl_image_frame),
|
||||
result_image, self.vdl_image_step)
|
||||
self.vdl_image_step += 1
|
||||
# each frame can display ten pictures at most.
|
||||
if self.vdl_image_step % 10 == 0:
|
||||
self.vdl_image_step = 0
|
||||
self.vdl_image_frame += 1
|
||||
|
||||
def on_epoch_end(self, status):
|
||||
mode = status['mode']
|
||||
if dist.get_world_size() < 2 or dist.get_rank() == 0:
|
||||
if mode == 'eval':
|
||||
for metric in self.model._metrics:
|
||||
for key, map_value in metric.get_results().items():
|
||||
self.vdl_writer.add_scalar("{}-mAP".format(key),
|
||||
map_value[0],
|
||||
self.vdl_mAP_step)
|
||||
self.vdl_mAP_step += 1
|
||||
|
||||
|
||||
class WandbCallback(Callback):
|
||||
def __init__(self, model):
|
||||
super(WandbCallback, self).__init__(model)
|
||||
|
||||
try:
|
||||
import wandb
|
||||
self.wandb = wandb
|
||||
except Exception as e:
|
||||
logger.error('wandb not found, please install wandb. '
|
||||
'Use: `pip install wandb`.')
|
||||
raise e
|
||||
|
||||
self.wandb_params = model.cfg.get('wandb', None)
|
||||
self.save_dir = os.path.join(self.model.cfg.save_dir,
|
||||
self.model.cfg.filename)
|
||||
if self.wandb_params is None:
|
||||
self.wandb_params = {}
|
||||
for k, v in model.cfg.items():
|
||||
if k.startswith("wandb_"):
|
||||
self.wandb_params.update({k.lstrip("wandb_"): v})
|
||||
|
||||
self._run = None
|
||||
if dist.get_world_size() < 2 or dist.get_rank() == 0:
|
||||
_ = self.run
|
||||
self.run.config.update(self.model.cfg)
|
||||
self.run.define_metric("epoch")
|
||||
self.run.define_metric("eval/*", step_metric="epoch")
|
||||
|
||||
self.best_ap = -1000.
|
||||
self.fps = []
|
||||
|
||||
@property
|
||||
def run(self):
|
||||
if self._run is None:
|
||||
if self.wandb.run is not None:
|
||||
logger.info(
|
||||
"There is an ongoing wandb run which will be used"
|
||||
"for logging. Please use `wandb.finish()` to end that"
|
||||
"if the behaviour is not intended")
|
||||
self._run = self.wandb.run
|
||||
else:
|
||||
self._run = self.wandb.init(**self.wandb_params)
|
||||
return self._run
|
||||
|
||||
def save_model(self,
|
||||
optimizer,
|
||||
save_dir,
|
||||
save_name,
|
||||
last_epoch,
|
||||
ema_model=None,
|
||||
ap=None,
|
||||
fps=None,
|
||||
tags=None):
|
||||
if dist.get_world_size() < 2 or dist.get_rank() == 0:
|
||||
model_path = os.path.join(save_dir, save_name)
|
||||
metadata = {}
|
||||
metadata["last_epoch"] = last_epoch
|
||||
if ap:
|
||||
metadata["ap"] = ap
|
||||
|
||||
if fps:
|
||||
metadata["fps"] = fps
|
||||
|
||||
if ema_model is None:
|
||||
ema_artifact = self.wandb.Artifact(
|
||||
name="ema_model-{}".format(self.run.id),
|
||||
type="model",
|
||||
metadata=metadata)
|
||||
model_artifact = self.wandb.Artifact(
|
||||
name="model-{}".format(self.run.id),
|
||||
type="model",
|
||||
metadata=metadata)
|
||||
|
||||
ema_artifact.add_file(model_path + ".pdema", name="model_ema")
|
||||
model_artifact.add_file(model_path + ".pdparams", name="model")
|
||||
|
||||
self.run.log_artifact(ema_artifact, aliases=tags)
|
||||
self.run.log_artfact(model_artifact, aliases=tags)
|
||||
else:
|
||||
model_artifact = self.wandb.Artifact(
|
||||
name="model-{}".format(self.run.id),
|
||||
type="model",
|
||||
metadata=metadata)
|
||||
model_artifact.add_file(model_path + ".pdparams", name="model")
|
||||
self.run.log_artifact(model_artifact, aliases=tags)
|
||||
|
||||
def on_step_end(self, status):
|
||||
|
||||
mode = status['mode']
|
||||
if dist.get_world_size() < 2 or dist.get_rank() == 0:
|
||||
if mode == 'train':
|
||||
training_status = status['training_status'].get()
|
||||
for k, v in training_status.items():
|
||||
training_status[k] = float(v)
|
||||
|
||||
# calculate ips, data_cost, batch_cost
|
||||
batch_time = status['batch_time']
|
||||
data_time = status['data_time']
|
||||
batch_size = self.model.cfg['{}Reader'.format(mode.capitalize(
|
||||
))]['batch_size']
|
||||
|
||||
ips = float(batch_size) / float(batch_time.avg)
|
||||
data_cost = float(data_time.avg)
|
||||
batch_cost = float(batch_time.avg)
|
||||
|
||||
metrics = {"train/" + k: v for k, v in training_status.items()}
|
||||
|
||||
metrics["train/ips"] = ips
|
||||
metrics["train/data_cost"] = data_cost
|
||||
metrics["train/batch_cost"] = batch_cost
|
||||
|
||||
self.fps.append(ips)
|
||||
self.run.log(metrics)
|
||||
|
||||
def on_epoch_end(self, status):
|
||||
mode = status['mode']
|
||||
epoch_id = status['epoch_id']
|
||||
save_name = None
|
||||
if dist.get_world_size() < 2 or dist.get_rank() == 0:
|
||||
if mode == 'train':
|
||||
fps = sum(self.fps) / len(self.fps)
|
||||
self.fps = []
|
||||
|
||||
end_epoch = self.model.cfg.epoch
|
||||
if (
|
||||
epoch_id + 1
|
||||
) % self.model.cfg.snapshot_epoch == 0 or epoch_id == end_epoch - 1:
|
||||
save_name = str(
|
||||
epoch_id) if epoch_id != end_epoch - 1 else "model_final"
|
||||
tags = ["latest", "epoch_{}".format(epoch_id)]
|
||||
self.save_model(
|
||||
self.model.optimizer,
|
||||
self.save_dir,
|
||||
save_name,
|
||||
epoch_id + 1,
|
||||
self.model.use_ema,
|
||||
fps=fps,
|
||||
tags=tags)
|
||||
if mode == 'eval':
|
||||
sample_num = status['sample_num']
|
||||
cost_time = status['cost_time']
|
||||
|
||||
fps = sample_num / cost_time
|
||||
|
||||
merged_dict = {}
|
||||
for metric in self.model._metrics:
|
||||
for key, map_value in metric.get_results().items():
|
||||
merged_dict["eval/{}-mAP".format(key)] = map_value[0]
|
||||
merged_dict["epoch"] = status["epoch_id"]
|
||||
merged_dict["eval/fps"] = sample_num / cost_time
|
||||
|
||||
self.run.log(merged_dict)
|
||||
|
||||
if 'save_best_model' in status and status['save_best_model']:
|
||||
for metric in self.model._metrics:
|
||||
map_res = metric.get_results()
|
||||
if 'pose3d' in map_res:
|
||||
key = 'pose3d'
|
||||
elif 'bbox' in map_res:
|
||||
key = 'bbox'
|
||||
elif 'keypoint' in map_res:
|
||||
key = 'keypoint'
|
||||
else:
|
||||
key = 'mask'
|
||||
if key not in map_res:
|
||||
logger.warning("Evaluation results empty, this may be due to " \
|
||||
"training iterations being too few or not " \
|
||||
"loading the correct weights.")
|
||||
return
|
||||
if map_res[key][0] >= self.best_ap:
|
||||
self.best_ap = map_res[key][0]
|
||||
save_name = 'best_model'
|
||||
tags = ["best", "epoch_{}".format(epoch_id)]
|
||||
|
||||
self.save_model(
|
||||
self.model.optimizer,
|
||||
self.save_dir,
|
||||
save_name,
|
||||
last_epoch=epoch_id + 1,
|
||||
ema_model=self.model.use_ema,
|
||||
ap=abs(self.best_ap),
|
||||
fps=fps,
|
||||
tags=tags)
|
||||
|
||||
def on_train_end(self, status):
|
||||
self.run.finish()
|
||||
|
||||
|
||||
class SniperProposalsGenerator(Callback):
|
||||
def __init__(self, model):
|
||||
super(SniperProposalsGenerator, self).__init__(model)
|
||||
ori_dataset = self.model.dataset
|
||||
self.dataset = self._create_new_dataset(ori_dataset)
|
||||
self.loader = self.model.loader
|
||||
self.cfg = self.model.cfg
|
||||
self.infer_model = self.model.model
|
||||
|
||||
def _create_new_dataset(self, ori_dataset):
|
||||
dataset = copy.deepcopy(ori_dataset)
|
||||
# init anno_cropper
|
||||
dataset.init_anno_cropper()
|
||||
# generate infer roidbs
|
||||
ori_roidbs = dataset.get_ori_roidbs()
|
||||
roidbs = dataset.anno_cropper.crop_infer_anno_records(ori_roidbs)
|
||||
# set new roidbs
|
||||
dataset.set_roidbs(roidbs)
|
||||
|
||||
return dataset
|
||||
|
||||
def _eval_with_loader(self, loader):
|
||||
results = []
|
||||
with paddle.no_grad():
|
||||
self.infer_model.eval()
|
||||
for step_id, data in enumerate(loader):
|
||||
outs = self.infer_model(data)
|
||||
for key in ['im_shape', 'scale_factor', 'im_id']:
|
||||
outs[key] = data[key]
|
||||
for key, value in outs.items():
|
||||
if hasattr(value, 'numpy'):
|
||||
outs[key] = value.numpy()
|
||||
|
||||
results.append(outs)
|
||||
|
||||
return results
|
||||
|
||||
def on_train_end(self, status):
|
||||
self.loader.dataset = self.dataset
|
||||
results = self._eval_with_loader(self.loader)
|
||||
results = self.dataset.anno_cropper.aggregate_chips_detections(results)
|
||||
# sniper
|
||||
proposals = []
|
||||
clsid2catid = {v: k for k, v in self.dataset.catid2clsid.items()}
|
||||
for outs in results:
|
||||
batch_res = get_infer_results(outs, clsid2catid)
|
||||
start = 0
|
||||
for i, im_id in enumerate(outs['im_id']):
|
||||
bbox_num = outs['bbox_num']
|
||||
end = start + bbox_num[i]
|
||||
bbox_res = batch_res['bbox'][start:end] \
|
||||
if 'bbox' in batch_res else None
|
||||
if bbox_res:
|
||||
proposals += bbox_res
|
||||
logger.info("save proposals in {}".format(self.cfg.proposals_path))
|
||||
with open(self.cfg.proposals_path, 'w') as f:
|
||||
json.dump(proposals, f)
|
||||
50
rtdetr_paddle/ppdet/engine/env.py
Normal file
50
rtdetr_paddle/ppdet/engine/env.py
Normal file
@@ -0,0 +1,50 @@
|
||||
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import os
|
||||
import random
|
||||
import numpy as np
|
||||
|
||||
import paddle
|
||||
from paddle.distributed import fleet
|
||||
|
||||
__all__ = ['init_parallel_env', 'set_random_seed', 'init_fleet_env']
|
||||
|
||||
|
||||
def init_fleet_env(find_unused_parameters=False):
|
||||
strategy = fleet.DistributedStrategy()
|
||||
strategy.find_unused_parameters = find_unused_parameters
|
||||
fleet.init(is_collective=True, strategy=strategy)
|
||||
|
||||
|
||||
def init_parallel_env():
|
||||
env = os.environ
|
||||
dist = 'PADDLE_TRAINER_ID' in env and 'PADDLE_TRAINERS_NUM' in env
|
||||
if dist:
|
||||
trainer_id = int(env['PADDLE_TRAINER_ID'])
|
||||
local_seed = (99 + trainer_id)
|
||||
random.seed(local_seed)
|
||||
np.random.seed(local_seed)
|
||||
|
||||
paddle.distributed.init_parallel_env()
|
||||
|
||||
|
||||
def set_random_seed(seed):
|
||||
paddle.seed(seed)
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
349
rtdetr_paddle/ppdet/engine/export_utils.py
Normal file
349
rtdetr_paddle/ppdet/engine/export_utils.py
Normal file
@@ -0,0 +1,349 @@
|
||||
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import os
|
||||
import yaml
|
||||
from collections import OrderedDict
|
||||
|
||||
import paddle
|
||||
from ppdet.data.source.category import get_categories
|
||||
|
||||
from ppdet.utils.logger import setup_logger
|
||||
logger = setup_logger('ppdet.engine')
|
||||
|
||||
# Global dictionary
|
||||
TRT_MIN_SUBGRAPH = {
|
||||
'YOLO': 3,
|
||||
'PPYOLOE': 3,
|
||||
'SSD': 60,
|
||||
'RCNN': 40,
|
||||
'RetinaNet': 40,
|
||||
'S2ANet': 80,
|
||||
'EfficientDet': 40,
|
||||
'Face': 3,
|
||||
'TTFNet': 60,
|
||||
'FCOS': 16,
|
||||
'SOLOv2': 60,
|
||||
'HigherHRNet': 3,
|
||||
'HRNet': 3,
|
||||
'DeepSORT': 3,
|
||||
'ByteTrack': 10,
|
||||
'CenterTrack': 5,
|
||||
'JDE': 10,
|
||||
'FairMOT': 5,
|
||||
'GFL': 16,
|
||||
'PicoDet': 3,
|
||||
'CenterNet': 5,
|
||||
'TOOD': 5,
|
||||
'YOLOX': 8,
|
||||
'YOLOF': 40,
|
||||
'METRO_Body': 3,
|
||||
'DETR': 3,
|
||||
}
|
||||
|
||||
KEYPOINT_ARCH = ['HigherHRNet', 'TopDownHRNet']
|
||||
MOT_ARCH = ['JDE', 'FairMOT', 'DeepSORT', 'ByteTrack', 'CenterTrack']
|
||||
|
||||
TO_STATIC_SPEC = {
|
||||
'yolov3_darknet53_270e_coco': [{
|
||||
'im_id': paddle.static.InputSpec(
|
||||
name='im_id', shape=[-1, 1], dtype='float32'),
|
||||
'is_crowd': paddle.static.InputSpec(
|
||||
name='is_crowd', shape=[-1, 50], dtype='float32'),
|
||||
'gt_bbox': paddle.static.InputSpec(
|
||||
name='gt_bbox', shape=[-1, 50, 4], dtype='float32'),
|
||||
'curr_iter': paddle.static.InputSpec(
|
||||
name='curr_iter', shape=[-1], dtype='float32'),
|
||||
'image': paddle.static.InputSpec(
|
||||
name='image', shape=[-1, 3, -1, -1], dtype='float32'),
|
||||
'im_shape': paddle.static.InputSpec(
|
||||
name='im_shape', shape=[-1, 2], dtype='float32'),
|
||||
'scale_factor': paddle.static.InputSpec(
|
||||
name='scale_factor', shape=[-1, 2], dtype='float32'),
|
||||
'target0': paddle.static.InputSpec(
|
||||
name='target0', shape=[-1, 3, 86, -1, -1], dtype='float32'),
|
||||
'target1': paddle.static.InputSpec(
|
||||
name='target1', shape=[-1, 3, 86, -1, -1], dtype='float32'),
|
||||
'target2': paddle.static.InputSpec(
|
||||
name='target2', shape=[-1, 3, 86, -1, -1], dtype='float32'),
|
||||
}],
|
||||
'tinypose_128x96': [{
|
||||
'center': paddle.static.InputSpec(
|
||||
name='center', shape=[-1, 2], dtype='float32'),
|
||||
'scale': paddle.static.InputSpec(
|
||||
name='scale', shape=[-1, 2], dtype='float32'),
|
||||
'im_id': paddle.static.InputSpec(
|
||||
name='im_id', shape=[-1, 1], dtype='float32'),
|
||||
'image': paddle.static.InputSpec(
|
||||
name='image', shape=[-1, 3, 128, 96], dtype='float32'),
|
||||
'score': paddle.static.InputSpec(
|
||||
name='score', shape=[-1], dtype='float32'),
|
||||
'rotate': paddle.static.InputSpec(
|
||||
name='rotate', shape=[-1], dtype='float32'),
|
||||
'target': paddle.static.InputSpec(
|
||||
name='target', shape=[-1, 17, 32, 24], dtype='float32'),
|
||||
'target_weight': paddle.static.InputSpec(
|
||||
name='target_weight', shape=[-1, 17, 1], dtype='float32'),
|
||||
}],
|
||||
'fcos_r50_fpn_1x_coco': [{
|
||||
'im_id': paddle.static.InputSpec(
|
||||
name='im_id', shape=[-1, 1], dtype='float32'),
|
||||
'curr_iter': paddle.static.InputSpec(
|
||||
name='curr_iter', shape=[-1], dtype='float32'),
|
||||
'image': paddle.static.InputSpec(
|
||||
name='image', shape=[-1, 3, -1, -1], dtype='float32'),
|
||||
'im_shape': paddle.static.InputSpec(
|
||||
name='im_shape', shape=[-1, 2], dtype='float32'),
|
||||
'scale_factor': paddle.static.InputSpec(
|
||||
name='scale_factor', shape=[-1, 2], dtype='float32'),
|
||||
'reg_target0': paddle.static.InputSpec(
|
||||
name='reg_target0', shape=[-1, 160, 160, 4], dtype='float32'),
|
||||
'labels0': paddle.static.InputSpec(
|
||||
name='labels0', shape=[-1, 160, 160, 1], dtype='int32'),
|
||||
'centerness0': paddle.static.InputSpec(
|
||||
name='centerness0', shape=[-1, 160, 160, 1], dtype='float32'),
|
||||
'reg_target1': paddle.static.InputSpec(
|
||||
name='reg_target1', shape=[-1, 80, 80, 4], dtype='float32'),
|
||||
'labels1': paddle.static.InputSpec(
|
||||
name='labels1', shape=[-1, 80, 80, 1], dtype='int32'),
|
||||
'centerness1': paddle.static.InputSpec(
|
||||
name='centerness1', shape=[-1, 80, 80, 1], dtype='float32'),
|
||||
'reg_target2': paddle.static.InputSpec(
|
||||
name='reg_target2', shape=[-1, 40, 40, 4], dtype='float32'),
|
||||
'labels2': paddle.static.InputSpec(
|
||||
name='labels2', shape=[-1, 40, 40, 1], dtype='int32'),
|
||||
'centerness2': paddle.static.InputSpec(
|
||||
name='centerness2', shape=[-1, 40, 40, 1], dtype='float32'),
|
||||
'reg_target3': paddle.static.InputSpec(
|
||||
name='reg_target3', shape=[-1, 20, 20, 4], dtype='float32'),
|
||||
'labels3': paddle.static.InputSpec(
|
||||
name='labels3', shape=[-1, 20, 20, 1], dtype='int32'),
|
||||
'centerness3': paddle.static.InputSpec(
|
||||
name='centerness3', shape=[-1, 20, 20, 1], dtype='float32'),
|
||||
'reg_target4': paddle.static.InputSpec(
|
||||
name='reg_target4', shape=[-1, 10, 10, 4], dtype='float32'),
|
||||
'labels4': paddle.static.InputSpec(
|
||||
name='labels4', shape=[-1, 10, 10, 1], dtype='int32'),
|
||||
'centerness4': paddle.static.InputSpec(
|
||||
name='centerness4', shape=[-1, 10, 10, 1], dtype='float32'),
|
||||
}],
|
||||
'picodet_s_320_coco_lcnet': [{
|
||||
'im_id': paddle.static.InputSpec(
|
||||
name='im_id', shape=[-1, 1], dtype='float32'),
|
||||
'is_crowd': paddle.static.InputSpec(
|
||||
name='is_crowd', shape=[-1, -1, 1], dtype='float32'),
|
||||
'gt_class': paddle.static.InputSpec(
|
||||
name='gt_class', shape=[-1, -1, 1], dtype='int32'),
|
||||
'gt_bbox': paddle.static.InputSpec(
|
||||
name='gt_bbox', shape=[-1, -1, 4], dtype='float32'),
|
||||
'curr_iter': paddle.static.InputSpec(
|
||||
name='curr_iter', shape=[-1], dtype='float32'),
|
||||
'image': paddle.static.InputSpec(
|
||||
name='image', shape=[-1, 3, -1, -1], dtype='float32'),
|
||||
'im_shape': paddle.static.InputSpec(
|
||||
name='im_shape', shape=[-1, 2], dtype='float32'),
|
||||
'scale_factor': paddle.static.InputSpec(
|
||||
name='scale_factor', shape=[-1, 2], dtype='float32'),
|
||||
'pad_gt_mask': paddle.static.InputSpec(
|
||||
name='pad_gt_mask', shape=[-1, -1, 1], dtype='float32'),
|
||||
}],
|
||||
'ppyoloe_crn_s_300e_coco': [{
|
||||
'im_id': paddle.static.InputSpec(
|
||||
name='im_id', shape=[-1, 1], dtype='float32'),
|
||||
'is_crowd': paddle.static.InputSpec(
|
||||
name='is_crowd', shape=[-1, -1, 1], dtype='float32'),
|
||||
'gt_class': paddle.static.InputSpec(
|
||||
name='gt_class', shape=[-1, -1, 1], dtype='int32'),
|
||||
'gt_bbox': paddle.static.InputSpec(
|
||||
name='gt_bbox', shape=[-1, -1, 4], dtype='float32'),
|
||||
'curr_iter': paddle.static.InputSpec(
|
||||
name='curr_iter', shape=[-1], dtype='float32'),
|
||||
'image': paddle.static.InputSpec(
|
||||
name='image', shape=[-1, 3, -1, -1], dtype='float32'),
|
||||
'im_shape': paddle.static.InputSpec(
|
||||
name='im_shape', shape=[-1, 2], dtype='float32'),
|
||||
'scale_factor': paddle.static.InputSpec(
|
||||
name='scale_factor', shape=[-1, 2], dtype='float32'),
|
||||
'pad_gt_mask': paddle.static.InputSpec(
|
||||
name='pad_gt_mask', shape=[-1, -1, 1], dtype='float32'),
|
||||
}],
|
||||
}
|
||||
|
||||
|
||||
def apply_to_static(config, model):
|
||||
filename = config.get('filename', None)
|
||||
spec = TO_STATIC_SPEC.get(filename, None)
|
||||
model = paddle.jit.to_static(model, input_spec=spec)
|
||||
logger.info("Successfully to apply @to_static with specs: {}".format(spec))
|
||||
return model
|
||||
|
||||
|
||||
def _prune_input_spec(input_spec, program, targets):
|
||||
# try to prune static program to figure out pruned input spec
|
||||
# so we perform following operations in static mode
|
||||
device = paddle.get_device()
|
||||
paddle.enable_static()
|
||||
paddle.set_device(device)
|
||||
pruned_input_spec = [{}]
|
||||
program = program.clone()
|
||||
program = program._prune(targets=targets)
|
||||
global_block = program.global_block()
|
||||
for name, spec in input_spec[0].items():
|
||||
try:
|
||||
v = global_block.var(name)
|
||||
pruned_input_spec[0][name] = spec
|
||||
except Exception:
|
||||
pass
|
||||
paddle.disable_static(place=device)
|
||||
return pruned_input_spec
|
||||
|
||||
|
||||
def _parse_reader(reader_cfg, dataset_cfg, metric, arch, image_shape):
|
||||
preprocess_list = []
|
||||
|
||||
anno_file = dataset_cfg.get_anno()
|
||||
|
||||
clsid2catid, catid2name = get_categories(metric, anno_file, arch)
|
||||
|
||||
label_list = [str(cat) for cat in catid2name.values()]
|
||||
|
||||
fuse_normalize = reader_cfg.get('fuse_normalize', False)
|
||||
sample_transforms = reader_cfg['sample_transforms']
|
||||
for st in sample_transforms[1:]:
|
||||
for key, value in st.items():
|
||||
p = {'type': key}
|
||||
if key == 'Resize':
|
||||
if int(image_shape[1]) != -1:
|
||||
value['target_size'] = image_shape[1:]
|
||||
value['interp'] = value.get('interp', 1) # cv2.INTER_LINEAR
|
||||
if fuse_normalize and key == 'NormalizeImage':
|
||||
continue
|
||||
p.update(value)
|
||||
preprocess_list.append(p)
|
||||
batch_transforms = reader_cfg.get('batch_transforms', None)
|
||||
if batch_transforms:
|
||||
for bt in batch_transforms:
|
||||
for key, value in bt.items():
|
||||
# for deploy/infer, use PadStride(stride) instead PadBatch(pad_to_stride)
|
||||
if key == 'PadBatch':
|
||||
preprocess_list.append({
|
||||
'type': 'PadStride',
|
||||
'stride': value['pad_to_stride']
|
||||
})
|
||||
break
|
||||
|
||||
return preprocess_list, label_list
|
||||
|
||||
|
||||
def _parse_tracker(tracker_cfg):
|
||||
tracker_params = {}
|
||||
for k, v in tracker_cfg.items():
|
||||
tracker_params.update({k: v})
|
||||
return tracker_params
|
||||
|
||||
|
||||
def _dump_infer_config(config, path, image_shape, model):
|
||||
arch_state = False
|
||||
from ppdet.core.config.yaml_helpers import setup_orderdict
|
||||
setup_orderdict()
|
||||
use_dynamic_shape = True if image_shape[2] == -1 else False
|
||||
infer_cfg = OrderedDict({
|
||||
'mode': 'paddle',
|
||||
'draw_threshold': 0.5,
|
||||
'metric': config['metric'],
|
||||
'use_dynamic_shape': use_dynamic_shape
|
||||
})
|
||||
export_onnx = config.get('export_onnx', False)
|
||||
export_eb = config.get('export_eb', False)
|
||||
|
||||
infer_arch = config['architecture']
|
||||
if 'RCNN' in infer_arch and export_onnx:
|
||||
logger.warning(
|
||||
"Exporting RCNN model to ONNX only support batch_size = 1")
|
||||
infer_cfg['export_onnx'] = True
|
||||
infer_cfg['export_eb'] = export_eb
|
||||
|
||||
if infer_arch in MOT_ARCH:
|
||||
if infer_arch == 'DeepSORT':
|
||||
tracker_cfg = config['DeepSORTTracker']
|
||||
elif infer_arch == 'CenterTrack':
|
||||
tracker_cfg = config['CenterTracker']
|
||||
else:
|
||||
tracker_cfg = config['JDETracker']
|
||||
infer_cfg['tracker'] = _parse_tracker(tracker_cfg)
|
||||
|
||||
for arch, min_subgraph_size in TRT_MIN_SUBGRAPH.items():
|
||||
if arch in infer_arch:
|
||||
infer_cfg['arch'] = arch
|
||||
infer_cfg['min_subgraph_size'] = min_subgraph_size
|
||||
arch_state = True
|
||||
break
|
||||
|
||||
if infer_arch == 'PPYOLOEWithAuxHead':
|
||||
infer_arch = 'PPYOLOE'
|
||||
|
||||
if infer_arch in ['PPYOLOE', 'YOLOX', 'YOLOF']:
|
||||
infer_cfg['arch'] = infer_arch
|
||||
infer_cfg['min_subgraph_size'] = TRT_MIN_SUBGRAPH[infer_arch]
|
||||
arch_state = True
|
||||
|
||||
if not arch_state:
|
||||
logger.error(
|
||||
'Architecture: {} is not supported for exporting model now.\n'.
|
||||
format(infer_arch) +
|
||||
'Please set TRT_MIN_SUBGRAPH in ppdet/engine/export_utils.py')
|
||||
os._exit(0)
|
||||
if 'mask_head' in config[config['architecture']] and config[config[
|
||||
'architecture']]['mask_head']:
|
||||
infer_cfg['mask'] = True
|
||||
label_arch = 'detection_arch'
|
||||
if infer_arch in KEYPOINT_ARCH:
|
||||
label_arch = 'keypoint_arch'
|
||||
|
||||
if infer_arch in MOT_ARCH:
|
||||
if config['metric'] in ['COCO', 'VOC']:
|
||||
# MOT model run as Detector
|
||||
reader_cfg = config['TestReader']
|
||||
dataset_cfg = config['TestDataset']
|
||||
else:
|
||||
# 'metric' in ['MOT', 'MCMOT', 'KITTI']
|
||||
label_arch = 'mot_arch'
|
||||
reader_cfg = config['TestMOTReader']
|
||||
dataset_cfg = config['TestMOTDataset']
|
||||
else:
|
||||
reader_cfg = config['TestReader']
|
||||
dataset_cfg = config['TestDataset']
|
||||
|
||||
infer_cfg['Preprocess'], infer_cfg['label_list'] = _parse_reader(
|
||||
reader_cfg, dataset_cfg, config['metric'], label_arch, image_shape[1:])
|
||||
|
||||
if infer_arch == 'PicoDet':
|
||||
if hasattr(config, 'export') and config['export'].get(
|
||||
'post_process',
|
||||
False) and not config['export'].get('benchmark', False):
|
||||
infer_cfg['arch'] = 'GFL'
|
||||
head_name = 'PicoHeadV2' if config['PicoHeadV2'] else 'PicoHead'
|
||||
infer_cfg['NMS'] = config[head_name]['nms']
|
||||
# In order to speed up the prediction, the threshold of nms
|
||||
# is adjusted here, which can be changed in infer_cfg.yml
|
||||
config[head_name]['nms']["score_threshold"] = 0.3
|
||||
config[head_name]['nms']["nms_threshold"] = 0.5
|
||||
infer_cfg['fpn_stride'] = config[head_name]['fpn_stride']
|
||||
|
||||
yaml.dump(infer_cfg, open(path, 'w'))
|
||||
logger.info("Export inference config file to {}".format(os.path.join(path)))
|
||||
966
rtdetr_paddle/ppdet/engine/trainer.py
Normal file
966
rtdetr_paddle/ppdet/engine/trainer.py
Normal file
@@ -0,0 +1,966 @@
|
||||
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import os
|
||||
import sys
|
||||
import copy
|
||||
import time
|
||||
from tqdm import tqdm
|
||||
|
||||
import numpy as np
|
||||
import typing
|
||||
from PIL import Image, ImageOps, ImageFile
|
||||
|
||||
ImageFile.LOAD_TRUNCATED_IMAGES = True
|
||||
|
||||
import paddle
|
||||
import paddle.nn as nn
|
||||
import paddle.distributed as dist
|
||||
from paddle.distributed import fleet
|
||||
from paddle.static import InputSpec
|
||||
from ppdet.optimizer import ModelEMA
|
||||
|
||||
from ppdet.core.workspace import create
|
||||
from ppdet.utils.checkpoint import load_weight, load_pretrain_weight
|
||||
from ppdet.utils.visualizer import visualize_results, save_result
|
||||
from ppdet.metrics import Metric, COCOMetric, VOCMetric, get_infer_results
|
||||
from ppdet.data.source.category import get_categories
|
||||
import ppdet.utils.stats as stats
|
||||
from ppdet.utils.fuse_utils import fuse_conv_bn
|
||||
from ppdet.utils import profiler
|
||||
from ppdet.modeling.post_process import multiclass_nms
|
||||
|
||||
from .callbacks import Callback, ComposeCallback, LogPrinter, Checkpointer, VisualDLWriter, WandbCallback
|
||||
from .export_utils import _dump_infer_config, _prune_input_spec, apply_to_static
|
||||
|
||||
from paddle.distributed.fleet.utils.hybrid_parallel_util import fused_allreduce_gradients
|
||||
|
||||
from ppdet.utils.logger import setup_logger
|
||||
logger = setup_logger('ppdet.engine')
|
||||
|
||||
__all__ = ['Trainer']
|
||||
|
||||
class Trainer(object):
|
||||
def __init__(self, cfg, mode='train'):
|
||||
self.cfg = cfg.copy()
|
||||
assert mode.lower() in ['train', 'eval', 'test'], \
|
||||
"mode should be 'train', 'eval' or 'test'"
|
||||
self.mode = mode.lower()
|
||||
self.optimizer = None
|
||||
self.is_loaded_weights = False
|
||||
self.use_amp = self.cfg.get('amp', False)
|
||||
self.amp_level = self.cfg.get('amp_level', 'O1')
|
||||
self.custom_white_list = self.cfg.get('custom_white_list', None)
|
||||
self.custom_black_list = self.cfg.get('custom_black_list', None)
|
||||
|
||||
# build data loader
|
||||
capital_mode = self.mode.capitalize()
|
||||
self.dataset = self.cfg['{}Dataset'.format(capital_mode)] = create(
|
||||
'{}Dataset'.format(capital_mode))()
|
||||
|
||||
if self.mode == 'train':
|
||||
self.loader = create('{}Reader'.format(capital_mode))(
|
||||
self.dataset, cfg.worker_num)
|
||||
|
||||
# build model
|
||||
if 'model' not in self.cfg:
|
||||
self.model = create(cfg.architecture)
|
||||
else:
|
||||
self.model = self.cfg.model
|
||||
self.is_loaded_weights = True
|
||||
|
||||
# EvalDataset build with BatchSampler to evaluate in single device
|
||||
# TODO: multi-device evaluate
|
||||
if self.mode == 'eval':
|
||||
self._eval_batch_sampler = paddle.io.BatchSampler(
|
||||
self.dataset, batch_size=self.cfg.EvalReader['batch_size'])
|
||||
reader_name = '{}Reader'.format(self.mode.capitalize())
|
||||
# If metric is VOC, need to be set collate_batch=False.
|
||||
if cfg.metric == 'VOC':
|
||||
self.cfg[reader_name]['collate_batch'] = False
|
||||
self.loader = create(reader_name)(self.dataset, cfg.worker_num,
|
||||
self._eval_batch_sampler)
|
||||
|
||||
# TestDataset build after user set images, skip loader creation here
|
||||
|
||||
# get Params
|
||||
print_params = self.cfg.get('print_params', False)
|
||||
if print_params:
|
||||
params = sum([
|
||||
p.numel() for n, p in self.model.named_parameters()
|
||||
if all([x not in n for x in ['_mean', '_variance', 'aux_']])
|
||||
]) # exclude BatchNorm running status
|
||||
logger.info('Model Params : {} M.'.format((params / 1e6).numpy()[
|
||||
0]))
|
||||
|
||||
# build optimizer in train mode
|
||||
if self.mode == 'train':
|
||||
steps_per_epoch = len(self.loader)
|
||||
if steps_per_epoch < 1:
|
||||
logger.warning(
|
||||
"Samples in dataset are less than batch_size, please set smaller batch_size in TrainReader."
|
||||
)
|
||||
self.lr = create('LearningRate')(steps_per_epoch)
|
||||
self.optimizer = create('OptimizerBuilder')(self.lr, self.model)
|
||||
|
||||
if self.use_amp and self.amp_level == 'O2':
|
||||
self.model, self.optimizer = paddle.amp.decorate(
|
||||
models=self.model,
|
||||
optimizers=self.optimizer,
|
||||
level=self.amp_level)
|
||||
self.use_ema = ('use_ema' in cfg and cfg['use_ema'])
|
||||
if self.use_ema:
|
||||
ema_decay = self.cfg.get('ema_decay', 0.9998)
|
||||
ema_decay_type = self.cfg.get('ema_decay_type', 'threshold')
|
||||
cycle_epoch = self.cfg.get('cycle_epoch', -1)
|
||||
ema_black_list = self.cfg.get('ema_black_list', None)
|
||||
ema_filter_no_grad = self.cfg.get('ema_filter_no_grad', False)
|
||||
self.ema = ModelEMA(
|
||||
self.model,
|
||||
decay=ema_decay,
|
||||
ema_decay_type=ema_decay_type,
|
||||
cycle_epoch=cycle_epoch,
|
||||
ema_black_list=ema_black_list,
|
||||
ema_filter_no_grad=ema_filter_no_grad)
|
||||
|
||||
self._nranks = dist.get_world_size()
|
||||
self._local_rank = dist.get_rank()
|
||||
|
||||
self.status = {}
|
||||
|
||||
self.start_epoch = 0
|
||||
self.end_epoch = 0 if 'epoch' not in cfg else cfg.epoch
|
||||
|
||||
# initial default callbacks
|
||||
self._init_callbacks()
|
||||
|
||||
# initial default metrics
|
||||
self._init_metrics()
|
||||
self._reset_metrics()
|
||||
|
||||
def _init_callbacks(self):
|
||||
if self.mode == 'train':
|
||||
self._callbacks = [LogPrinter(self), Checkpointer(self)]
|
||||
if self.cfg.get('use_vdl', False):
|
||||
self._callbacks.append(VisualDLWriter(self))
|
||||
if self.cfg.get('use_wandb', False) or 'wandb' in self.cfg:
|
||||
self._callbacks.append(WandbCallback(self))
|
||||
self._compose_callback = ComposeCallback(self._callbacks)
|
||||
elif self.mode == 'eval':
|
||||
self._callbacks = [LogPrinter(self)]
|
||||
self._compose_callback = ComposeCallback(self._callbacks)
|
||||
elif self.mode == 'test' and self.cfg.get('use_vdl', False):
|
||||
self._callbacks = [VisualDLWriter(self)]
|
||||
self._compose_callback = ComposeCallback(self._callbacks)
|
||||
else:
|
||||
self._callbacks = []
|
||||
self._compose_callback = None
|
||||
|
||||
def _init_metrics(self, validate=False):
|
||||
if self.mode == 'test' or (self.mode == 'train' and not validate):
|
||||
self._metrics = []
|
||||
return
|
||||
classwise = self.cfg['classwise'] if 'classwise' in self.cfg else False
|
||||
if self.cfg.metric == 'COCO':
|
||||
# TODO: bias should be unified
|
||||
bias = 1 if self.cfg.get('bias', False) else 0
|
||||
output_eval = self.cfg['output_eval'] \
|
||||
if 'output_eval' in self.cfg else None
|
||||
save_prediction_only = self.cfg.get('save_prediction_only', False)
|
||||
|
||||
# pass clsid2catid info to metric instance to avoid multiple loading
|
||||
# annotation file
|
||||
clsid2catid = {v: k for k, v in self.dataset.catid2clsid.items()} \
|
||||
if self.mode == 'eval' else None
|
||||
|
||||
# when do validation in train, annotation file should be get from
|
||||
# EvalReader instead of self.dataset(which is TrainReader)
|
||||
if self.mode == 'train' and validate:
|
||||
eval_dataset = self.cfg['EvalDataset']
|
||||
eval_dataset.check_or_download_dataset()
|
||||
anno_file = eval_dataset.get_anno()
|
||||
dataset = eval_dataset
|
||||
else:
|
||||
dataset = self.dataset
|
||||
anno_file = dataset.get_anno()
|
||||
|
||||
IouType = self.cfg['IouType'] if 'IouType' in self.cfg else 'bbox'
|
||||
self._metrics = [
|
||||
COCOMetric(
|
||||
anno_file=anno_file,
|
||||
clsid2catid=clsid2catid,
|
||||
classwise=classwise,
|
||||
output_eval=output_eval,
|
||||
bias=bias,
|
||||
IouType=IouType,
|
||||
save_prediction_only=save_prediction_only)
|
||||
]
|
||||
|
||||
elif self.cfg.metric == 'VOC':
|
||||
output_eval = self.cfg['output_eval'] \
|
||||
if 'output_eval' in self.cfg else None
|
||||
save_prediction_only = self.cfg.get('save_prediction_only', False)
|
||||
self._metrics = [
|
||||
VOCMetric(
|
||||
label_list=self.dataset.get_label_list(),
|
||||
class_num=self.cfg.num_classes,
|
||||
map_type=self.cfg.map_type,
|
||||
classwise=classwise,
|
||||
output_eval=output_eval,
|
||||
save_prediction_only=save_prediction_only)
|
||||
]
|
||||
else:
|
||||
logger.warning("Metric not support for metric type {}".format(
|
||||
self.cfg.metric))
|
||||
self._metrics = []
|
||||
|
||||
def _reset_metrics(self):
|
||||
for metric in self._metrics:
|
||||
metric.reset()
|
||||
|
||||
def register_callbacks(self, callbacks):
|
||||
callbacks = [c for c in list(callbacks) if c is not None]
|
||||
for c in callbacks:
|
||||
assert isinstance(c, Callback), \
|
||||
"metrics shoule be instances of subclass of Metric"
|
||||
self._callbacks.extend(callbacks)
|
||||
self._compose_callback = ComposeCallback(self._callbacks)
|
||||
|
||||
def register_metrics(self, metrics):
|
||||
metrics = [m for m in list(metrics) if m is not None]
|
||||
for m in metrics:
|
||||
assert isinstance(m, Metric), \
|
||||
"metrics shoule be instances of subclass of Metric"
|
||||
self._metrics.extend(metrics)
|
||||
|
||||
def load_weights(self, weights, ARSL_eval=False):
|
||||
if self.is_loaded_weights:
|
||||
return
|
||||
self.start_epoch = 0
|
||||
load_pretrain_weight(self.model, weights, ARSL_eval)
|
||||
logger.debug("Load weights {} to start training".format(weights))
|
||||
|
||||
def resume_weights(self, weights):
|
||||
self.start_epoch = load_weight(self.model, weights, self.optimizer,
|
||||
self.ema if self.use_ema else None)
|
||||
logger.debug("Resume weights of epoch {}".format(self.start_epoch))
|
||||
|
||||
def train(self, validate=False):
|
||||
assert self.mode == 'train', "Model not in 'train' mode"
|
||||
Init_mark = False
|
||||
if validate:
|
||||
self.cfg['EvalDataset'] = self.cfg.EvalDataset = create(
|
||||
"EvalDataset")()
|
||||
|
||||
model = self.model
|
||||
if self.cfg.get('to_static', False):
|
||||
model = apply_to_static(self.cfg, model)
|
||||
sync_bn = (getattr(self.cfg, 'norm_type', None) == 'sync_bn' and
|
||||
(self.cfg.use_gpu or self.cfg.use_mlu) and self._nranks > 1)
|
||||
if sync_bn:
|
||||
model = paddle.nn.SyncBatchNorm.convert_sync_batchnorm(model)
|
||||
|
||||
# enabel auto mixed precision mode
|
||||
if self.use_amp:
|
||||
scaler = paddle.amp.GradScaler(
|
||||
enable=self.cfg.use_gpu or self.cfg.use_npu or self.cfg.use_mlu,
|
||||
init_loss_scaling=self.cfg.get('init_loss_scaling', 1024))
|
||||
# get distributed model
|
||||
if self.cfg.get('fleet', False):
|
||||
model = fleet.distributed_model(model)
|
||||
self.optimizer = fleet.distributed_optimizer(self.optimizer)
|
||||
elif self._nranks > 1:
|
||||
find_unused_parameters = self.cfg[
|
||||
'find_unused_parameters'] if 'find_unused_parameters' in self.cfg else False
|
||||
model = paddle.DataParallel(
|
||||
model, find_unused_parameters=find_unused_parameters)
|
||||
|
||||
self.status.update({
|
||||
'epoch_id': self.start_epoch,
|
||||
'step_id': 0,
|
||||
'steps_per_epoch': len(self.loader)
|
||||
})
|
||||
|
||||
self.status['batch_time'] = stats.SmoothedValue(
|
||||
self.cfg.log_iter, fmt='{avg:.4f}')
|
||||
self.status['data_time'] = stats.SmoothedValue(
|
||||
self.cfg.log_iter, fmt='{avg:.4f}')
|
||||
self.status['training_status'] = stats.TrainingStats(self.cfg.log_iter)
|
||||
|
||||
profiler_options = self.cfg.get('profiler_options', None)
|
||||
|
||||
self._compose_callback.on_train_begin(self.status)
|
||||
|
||||
use_fused_allreduce_gradients = self.cfg[
|
||||
'use_fused_allreduce_gradients'] if 'use_fused_allreduce_gradients' in self.cfg else False
|
||||
|
||||
for epoch_id in range(self.start_epoch, self.cfg.epoch):
|
||||
self.status['mode'] = 'train'
|
||||
self.status['epoch_id'] = epoch_id
|
||||
self._compose_callback.on_epoch_begin(self.status)
|
||||
self.loader.dataset.set_epoch(epoch_id)
|
||||
model.train()
|
||||
iter_tic = time.time()
|
||||
for step_id, data in enumerate(self.loader):
|
||||
self.status['data_time'].update(time.time() - iter_tic)
|
||||
self.status['step_id'] = step_id
|
||||
profiler.add_profiler_step(profiler_options)
|
||||
self._compose_callback.on_step_begin(self.status)
|
||||
data['epoch_id'] = epoch_id
|
||||
if self.cfg.get('to_static',
|
||||
False) and 'image_file' in data.keys():
|
||||
data.pop('image_file')
|
||||
|
||||
if self.use_amp:
|
||||
if isinstance(
|
||||
model, paddle.
|
||||
DataParallel) and use_fused_allreduce_gradients:
|
||||
with model.no_sync():
|
||||
with paddle.amp.auto_cast(
|
||||
enable=self.cfg.use_gpu or
|
||||
self.cfg.use_npu or self.cfg.use_mlu,
|
||||
custom_white_list=self.custom_white_list,
|
||||
custom_black_list=self.custom_black_list,
|
||||
level=self.amp_level):
|
||||
# model forward
|
||||
outputs = model(data)
|
||||
loss = outputs['loss']
|
||||
# model backward
|
||||
scaled_loss = scaler.scale(loss)
|
||||
scaled_loss.backward()
|
||||
fused_allreduce_gradients(
|
||||
list(model.parameters()), None)
|
||||
else:
|
||||
with paddle.amp.auto_cast(
|
||||
enable=self.cfg.use_gpu or self.cfg.use_npu or
|
||||
self.cfg.use_mlu,
|
||||
custom_white_list=self.custom_white_list,
|
||||
custom_black_list=self.custom_black_list,
|
||||
level=self.amp_level):
|
||||
# model forward
|
||||
outputs = model(data)
|
||||
loss = outputs['loss']
|
||||
# model backward
|
||||
scaled_loss = scaler.scale(loss)
|
||||
scaled_loss.backward()
|
||||
# in dygraph mode, optimizer.minimize is equal to optimizer.step
|
||||
scaler.minimize(self.optimizer, scaled_loss)
|
||||
else:
|
||||
if isinstance(
|
||||
model, paddle.
|
||||
DataParallel) and use_fused_allreduce_gradients:
|
||||
with model.no_sync():
|
||||
# model forward
|
||||
outputs = model(data)
|
||||
loss = outputs['loss']
|
||||
# model backward
|
||||
loss.backward()
|
||||
fused_allreduce_gradients(
|
||||
list(model.parameters()), None)
|
||||
else:
|
||||
# model forward
|
||||
outputs = model(data)
|
||||
loss = outputs['loss']
|
||||
# model backward
|
||||
loss.backward()
|
||||
self.optimizer.step()
|
||||
curr_lr = self.optimizer.get_lr()
|
||||
self.lr.step()
|
||||
self.optimizer.clear_grad()
|
||||
self.status['learning_rate'] = curr_lr
|
||||
|
||||
if self._nranks < 2 or self._local_rank == 0:
|
||||
self.status['training_status'].update(outputs)
|
||||
|
||||
self.status['batch_time'].update(time.time() - iter_tic)
|
||||
self._compose_callback.on_step_end(self.status)
|
||||
if self.use_ema:
|
||||
self.ema.update()
|
||||
iter_tic = time.time()
|
||||
|
||||
is_snapshot = (self._nranks < 2 or (self._local_rank == 0 or self.cfg.metric == "Pose3DEval")) \
|
||||
and ((epoch_id + 1) % self.cfg.snapshot_epoch == 0 or epoch_id == self.end_epoch - 1)
|
||||
if is_snapshot and self.use_ema:
|
||||
# apply ema weight on model
|
||||
weight = copy.deepcopy(self.model.state_dict())
|
||||
self.model.set_dict(self.ema.apply())
|
||||
self.status['weight'] = weight
|
||||
|
||||
self._compose_callback.on_epoch_end(self.status)
|
||||
|
||||
if validate and is_snapshot:
|
||||
if not hasattr(self, '_eval_loader'):
|
||||
# build evaluation dataset and loader
|
||||
self._eval_dataset = self.cfg.EvalDataset
|
||||
self._eval_batch_sampler = \
|
||||
paddle.io.BatchSampler(
|
||||
self._eval_dataset,
|
||||
batch_size=self.cfg.EvalReader['batch_size'])
|
||||
# If metric is VOC, need to be set collate_batch=False.
|
||||
if self.cfg.metric == 'VOC':
|
||||
self.cfg['EvalReader']['collate_batch'] = False
|
||||
else:
|
||||
self._eval_loader = create('EvalReader')(
|
||||
self._eval_dataset,
|
||||
self.cfg.worker_num,
|
||||
batch_sampler=self._eval_batch_sampler)
|
||||
# if validation in training is enabled, metrics should be re-init
|
||||
# Init_mark makes sure this code will only execute once
|
||||
if validate and Init_mark == False:
|
||||
Init_mark = True
|
||||
self._init_metrics(validate=validate)
|
||||
self._reset_metrics()
|
||||
|
||||
with paddle.no_grad():
|
||||
self.status['save_best_model'] = True
|
||||
self._eval_with_loader(self._eval_loader)
|
||||
|
||||
if is_snapshot and self.use_ema:
|
||||
# reset original weight
|
||||
self.model.set_dict(weight)
|
||||
self.status.pop('weight')
|
||||
|
||||
self._compose_callback.on_train_end(self.status)
|
||||
|
||||
def _eval_with_loader(self, loader):
|
||||
sample_num = 0
|
||||
tic = time.time()
|
||||
self._compose_callback.on_epoch_begin(self.status)
|
||||
self.status['mode'] = 'eval'
|
||||
|
||||
self.model.eval()
|
||||
for step_id, data in enumerate(loader):
|
||||
self.status['step_id'] = step_id
|
||||
self._compose_callback.on_step_begin(self.status)
|
||||
# forward
|
||||
if self.use_amp:
|
||||
with paddle.amp.auto_cast(
|
||||
enable=self.cfg.use_gpu or self.cfg.use_npu or
|
||||
self.cfg.use_mlu,
|
||||
custom_white_list=self.custom_white_list,
|
||||
custom_black_list=self.custom_black_list,
|
||||
level=self.amp_level):
|
||||
outs = self.model(data)
|
||||
else:
|
||||
outs = self.model(data)
|
||||
|
||||
# update metrics
|
||||
for metric in self._metrics:
|
||||
metric.update(data, outs)
|
||||
|
||||
# multi-scale inputs: all inputs have same im_id
|
||||
if isinstance(data, typing.Sequence):
|
||||
sample_num += data[0]['im_id'].numpy().shape[0]
|
||||
else:
|
||||
sample_num += data['im_id'].numpy().shape[0]
|
||||
self._compose_callback.on_step_end(self.status)
|
||||
|
||||
self.status['sample_num'] = sample_num
|
||||
self.status['cost_time'] = time.time() - tic
|
||||
|
||||
# accumulate metric to log out
|
||||
for metric in self._metrics:
|
||||
metric.accumulate()
|
||||
metric.log()
|
||||
self._compose_callback.on_epoch_end(self.status)
|
||||
# reset metric states for metric may performed multiple times
|
||||
self._reset_metrics()
|
||||
|
||||
def evaluate(self):
|
||||
# get distributed model
|
||||
if self.cfg.get('fleet', False):
|
||||
self.model = fleet.distributed_model(self.model)
|
||||
self.optimizer = fleet.distributed_optimizer(self.optimizer)
|
||||
elif self._nranks > 1:
|
||||
find_unused_parameters = self.cfg[
|
||||
'find_unused_parameters'] if 'find_unused_parameters' in self.cfg else False
|
||||
self.model = paddle.DataParallel(
|
||||
self.model, find_unused_parameters=find_unused_parameters)
|
||||
with paddle.no_grad():
|
||||
self._eval_with_loader(self.loader)
|
||||
|
||||
def _eval_with_loader_slice(self,
|
||||
loader,
|
||||
slice_size=[640, 640],
|
||||
overlap_ratio=[0.25, 0.25],
|
||||
combine_method='nms',
|
||||
match_threshold=0.6,
|
||||
match_metric='iou'):
|
||||
sample_num = 0
|
||||
tic = time.time()
|
||||
self._compose_callback.on_epoch_begin(self.status)
|
||||
self.status['mode'] = 'eval'
|
||||
self.model.eval()
|
||||
merged_bboxs = []
|
||||
for step_id, data in enumerate(loader):
|
||||
self.status['step_id'] = step_id
|
||||
self._compose_callback.on_step_begin(self.status)
|
||||
# forward
|
||||
if self.use_amp:
|
||||
with paddle.amp.auto_cast(
|
||||
enable=self.cfg.use_gpu or self.cfg.use_npu or
|
||||
self.cfg.use_mlu,
|
||||
custom_white_list=self.custom_white_list,
|
||||
custom_black_list=self.custom_black_list,
|
||||
level=self.amp_level):
|
||||
outs = self.model(data)
|
||||
else:
|
||||
outs = self.model(data)
|
||||
|
||||
shift_amount = data['st_pix']
|
||||
outs['bbox'][:, 2:4] = outs['bbox'][:, 2:4] + shift_amount
|
||||
outs['bbox'][:, 4:6] = outs['bbox'][:, 4:6] + shift_amount
|
||||
merged_bboxs.append(outs['bbox'])
|
||||
|
||||
if data['is_last'] > 0:
|
||||
# merge matching predictions
|
||||
merged_results = {'bbox': []}
|
||||
if combine_method == 'nms':
|
||||
final_boxes = multiclass_nms(
|
||||
np.concatenate(merged_bboxs), self.cfg.num_classes,
|
||||
match_threshold, match_metric)
|
||||
merged_results['bbox'] = np.concatenate(final_boxes)
|
||||
elif combine_method == 'concat':
|
||||
merged_results['bbox'] = np.concatenate(merged_bboxs)
|
||||
else:
|
||||
raise ValueError(
|
||||
"Now only support 'nms' or 'concat' to fuse detection results."
|
||||
)
|
||||
merged_results['im_id'] = np.array([[0]])
|
||||
merged_results['bbox_num'] = np.array(
|
||||
[len(merged_results['bbox'])])
|
||||
|
||||
merged_bboxs = []
|
||||
data['im_id'] = data['ori_im_id']
|
||||
# update metrics
|
||||
for metric in self._metrics:
|
||||
metric.update(data, merged_results)
|
||||
|
||||
# multi-scale inputs: all inputs have same im_id
|
||||
if isinstance(data, typing.Sequence):
|
||||
sample_num += data[0]['im_id'].numpy().shape[0]
|
||||
else:
|
||||
sample_num += data['im_id'].numpy().shape[0]
|
||||
|
||||
self._compose_callback.on_step_end(self.status)
|
||||
|
||||
self.status['sample_num'] = sample_num
|
||||
self.status['cost_time'] = time.time() - tic
|
||||
|
||||
# accumulate metric to log out
|
||||
for metric in self._metrics:
|
||||
metric.accumulate()
|
||||
metric.log()
|
||||
self._compose_callback.on_epoch_end(self.status)
|
||||
# reset metric states for metric may performed multiple times
|
||||
self._reset_metrics()
|
||||
|
||||
def evaluate_slice(self,
|
||||
slice_size=[640, 640],
|
||||
overlap_ratio=[0.25, 0.25],
|
||||
combine_method='nms',
|
||||
match_threshold=0.6,
|
||||
match_metric='iou'):
|
||||
with paddle.no_grad():
|
||||
self._eval_with_loader_slice(self.loader, slice_size, overlap_ratio,
|
||||
combine_method, match_threshold,
|
||||
match_metric)
|
||||
|
||||
def slice_predict(self,
|
||||
images,
|
||||
slice_size=[640, 640],
|
||||
overlap_ratio=[0.25, 0.25],
|
||||
combine_method='nms',
|
||||
match_threshold=0.6,
|
||||
match_metric='iou',
|
||||
draw_threshold=0.5,
|
||||
output_dir='output',
|
||||
save_results=False,
|
||||
visualize=True):
|
||||
if not os.path.exists(output_dir):
|
||||
os.makedirs(output_dir)
|
||||
|
||||
self.dataset.set_slice_images(images, slice_size, overlap_ratio)
|
||||
loader = create('TestReader')(self.dataset, 0)
|
||||
imid2path = self.dataset.get_imid2path()
|
||||
|
||||
def setup_metrics_for_loader():
|
||||
# mem
|
||||
metrics = copy.deepcopy(self._metrics)
|
||||
mode = self.mode
|
||||
save_prediction_only = self.cfg[
|
||||
'save_prediction_only'] if 'save_prediction_only' in self.cfg else None
|
||||
output_eval = self.cfg[
|
||||
'output_eval'] if 'output_eval' in self.cfg else None
|
||||
|
||||
# modify
|
||||
self.mode = '_test'
|
||||
self.cfg['save_prediction_only'] = True
|
||||
self.cfg['output_eval'] = output_dir
|
||||
self.cfg['imid2path'] = imid2path
|
||||
self._init_metrics()
|
||||
|
||||
# restore
|
||||
self.mode = mode
|
||||
self.cfg.pop('save_prediction_only')
|
||||
if save_prediction_only is not None:
|
||||
self.cfg['save_prediction_only'] = save_prediction_only
|
||||
|
||||
self.cfg.pop('output_eval')
|
||||
if output_eval is not None:
|
||||
self.cfg['output_eval'] = output_eval
|
||||
|
||||
self.cfg.pop('imid2path')
|
||||
|
||||
_metrics = copy.deepcopy(self._metrics)
|
||||
self._metrics = metrics
|
||||
|
||||
return _metrics
|
||||
|
||||
if save_results:
|
||||
metrics = setup_metrics_for_loader()
|
||||
else:
|
||||
metrics = []
|
||||
|
||||
anno_file = self.dataset.get_anno()
|
||||
clsid2catid, catid2name = get_categories(
|
||||
self.cfg.metric, anno_file=anno_file)
|
||||
|
||||
# Run Infer
|
||||
self.status['mode'] = 'test'
|
||||
self.model.eval()
|
||||
|
||||
results = [] # all images
|
||||
merged_bboxs = [] # single image
|
||||
for step_id, data in enumerate(tqdm(loader)):
|
||||
self.status['step_id'] = step_id
|
||||
# forward
|
||||
with paddle.no_grad():
|
||||
outs = self.model(data)
|
||||
|
||||
outs['bbox'] = outs['bbox'].numpy() # only in test mode
|
||||
shift_amount = data['st_pix']
|
||||
outs['bbox'][:, 2:4] = outs['bbox'][:, 2:4] + shift_amount.numpy()
|
||||
outs['bbox'][:, 4:6] = outs['bbox'][:, 4:6] + shift_amount.numpy()
|
||||
merged_bboxs.append(outs['bbox'])
|
||||
|
||||
if data['is_last'] > 0:
|
||||
# merge matching predictions
|
||||
merged_results = {'bbox': []}
|
||||
if combine_method == 'nms':
|
||||
final_boxes = multiclass_nms(
|
||||
np.concatenate(merged_bboxs), self.cfg.num_classes,
|
||||
match_threshold, match_metric)
|
||||
merged_results['bbox'] = np.concatenate(final_boxes)
|
||||
elif combine_method == 'concat':
|
||||
merged_results['bbox'] = np.concatenate(merged_bboxs)
|
||||
else:
|
||||
raise ValueError(
|
||||
"Now only support 'nms' or 'concat' to fuse detection results."
|
||||
)
|
||||
merged_results['im_id'] = np.array([[0]])
|
||||
merged_results['bbox_num'] = np.array(
|
||||
[len(merged_results['bbox'])])
|
||||
|
||||
merged_bboxs = []
|
||||
data['im_id'] = data['ori_im_id']
|
||||
|
||||
for _m in metrics:
|
||||
_m.update(data, merged_results)
|
||||
|
||||
for key in ['im_shape', 'scale_factor', 'im_id']:
|
||||
if isinstance(data, typing.Sequence):
|
||||
merged_results[key] = data[0][key]
|
||||
else:
|
||||
merged_results[key] = data[key]
|
||||
for key, value in merged_results.items():
|
||||
if hasattr(value, 'numpy'):
|
||||
merged_results[key] = value.numpy()
|
||||
results.append(merged_results)
|
||||
|
||||
for _m in metrics:
|
||||
_m.accumulate()
|
||||
_m.reset()
|
||||
|
||||
if visualize:
|
||||
for outs in results:
|
||||
batch_res = get_infer_results(outs, clsid2catid)
|
||||
bbox_num = outs['bbox_num']
|
||||
|
||||
start = 0
|
||||
for i, im_id in enumerate(outs['im_id']):
|
||||
image_path = imid2path[int(im_id)]
|
||||
image = Image.open(image_path).convert('RGB')
|
||||
image = ImageOps.exif_transpose(image)
|
||||
self.status['original_image'] = np.array(image.copy())
|
||||
|
||||
end = start + bbox_num[i]
|
||||
bbox_res = batch_res['bbox'][start:end] \
|
||||
if 'bbox' in batch_res else None
|
||||
mask_res = batch_res['mask'][start:end] \
|
||||
if 'mask' in batch_res else None
|
||||
segm_res = batch_res['segm'][start:end] \
|
||||
if 'segm' in batch_res else None
|
||||
keypoint_res = batch_res['keypoint'][start:end] \
|
||||
if 'keypoint' in batch_res else None
|
||||
pose3d_res = batch_res['pose3d'][start:end] \
|
||||
if 'pose3d' in batch_res else None
|
||||
image = visualize_results(
|
||||
image, bbox_res, mask_res, segm_res, keypoint_res,
|
||||
pose3d_res, int(im_id), catid2name, draw_threshold)
|
||||
self.status['result_image'] = np.array(image.copy())
|
||||
if self._compose_callback:
|
||||
self._compose_callback.on_step_end(self.status)
|
||||
# save image with detection
|
||||
save_name = self._get_save_image_name(output_dir,
|
||||
image_path)
|
||||
logger.info("Detection bbox results save in {}".format(
|
||||
save_name))
|
||||
image.save(save_name, quality=95)
|
||||
|
||||
start = end
|
||||
|
||||
def predict(self,
|
||||
images,
|
||||
draw_threshold=0.5,
|
||||
output_dir='output',
|
||||
save_results=False,
|
||||
visualize=True):
|
||||
if not os.path.exists(output_dir):
|
||||
os.makedirs(output_dir)
|
||||
|
||||
self.dataset.set_images(images)
|
||||
loader = create('TestReader')(self.dataset, 0)
|
||||
|
||||
imid2path = self.dataset.get_imid2path()
|
||||
|
||||
def setup_metrics_for_loader():
|
||||
# mem
|
||||
metrics = copy.deepcopy(self._metrics)
|
||||
mode = self.mode
|
||||
save_prediction_only = self.cfg[
|
||||
'save_prediction_only'] if 'save_prediction_only' in self.cfg else None
|
||||
output_eval = self.cfg[
|
||||
'output_eval'] if 'output_eval' in self.cfg else None
|
||||
|
||||
# modify
|
||||
self.mode = '_test'
|
||||
self.cfg['save_prediction_only'] = True
|
||||
self.cfg['output_eval'] = output_dir
|
||||
self.cfg['imid2path'] = imid2path
|
||||
self._init_metrics()
|
||||
|
||||
# restore
|
||||
self.mode = mode
|
||||
self.cfg.pop('save_prediction_only')
|
||||
if save_prediction_only is not None:
|
||||
self.cfg['save_prediction_only'] = save_prediction_only
|
||||
|
||||
self.cfg.pop('output_eval')
|
||||
if output_eval is not None:
|
||||
self.cfg['output_eval'] = output_eval
|
||||
|
||||
self.cfg.pop('imid2path')
|
||||
|
||||
_metrics = copy.deepcopy(self._metrics)
|
||||
self._metrics = metrics
|
||||
|
||||
return _metrics
|
||||
|
||||
if save_results:
|
||||
metrics = setup_metrics_for_loader()
|
||||
else:
|
||||
metrics = []
|
||||
|
||||
anno_file = self.dataset.get_anno()
|
||||
clsid2catid, catid2name = get_categories(
|
||||
self.cfg.metric, anno_file=anno_file)
|
||||
|
||||
# Run Infer
|
||||
self.status['mode'] = 'test'
|
||||
self.model.eval()
|
||||
|
||||
results = []
|
||||
for step_id, data in enumerate(tqdm(loader)):
|
||||
self.status['step_id'] = step_id
|
||||
# forward
|
||||
with paddle.no_grad():
|
||||
if hasattr(self.model, 'modelTeacher'):
|
||||
outs = self.model.modelTeacher(data)
|
||||
else:
|
||||
outs = self.model(data)
|
||||
|
||||
for _m in metrics:
|
||||
_m.update(data, outs)
|
||||
|
||||
for key in ['im_shape', 'scale_factor', 'im_id']:
|
||||
if isinstance(data, typing.Sequence):
|
||||
outs[key] = data[0][key]
|
||||
else:
|
||||
outs[key] = data[key]
|
||||
for key, value in outs.items():
|
||||
if hasattr(value, 'numpy'):
|
||||
outs[key] = value.numpy()
|
||||
results.append(outs)
|
||||
|
||||
for _m in metrics:
|
||||
_m.accumulate()
|
||||
_m.reset()
|
||||
|
||||
if visualize:
|
||||
for outs in results:
|
||||
batch_res = get_infer_results(outs, clsid2catid)
|
||||
bbox_num = outs['bbox_num']
|
||||
|
||||
start = 0
|
||||
for i, im_id in enumerate(outs['im_id']):
|
||||
image_path = imid2path[int(im_id)]
|
||||
image = Image.open(image_path).convert('RGB')
|
||||
image = ImageOps.exif_transpose(image)
|
||||
self.status['original_image'] = np.array(image.copy())
|
||||
|
||||
end = start + bbox_num[i]
|
||||
bbox_res = batch_res['bbox'][start:end] \
|
||||
if 'bbox' in batch_res else None
|
||||
mask_res = batch_res['mask'][start:end] \
|
||||
if 'mask' in batch_res else None
|
||||
segm_res = batch_res['segm'][start:end] \
|
||||
if 'segm' in batch_res else None
|
||||
keypoint_res = batch_res['keypoint'][start:end] \
|
||||
if 'keypoint' in batch_res else None
|
||||
pose3d_res = batch_res['pose3d'][start:end] \
|
||||
if 'pose3d' in batch_res else None
|
||||
image = visualize_results(
|
||||
image, bbox_res, mask_res, segm_res, keypoint_res,
|
||||
pose3d_res, int(im_id), catid2name, draw_threshold)
|
||||
self.status['result_image'] = np.array(image.copy())
|
||||
if self._compose_callback:
|
||||
self._compose_callback.on_step_end(self.status)
|
||||
# save image with detection
|
||||
save_name = self._get_save_image_name(output_dir,
|
||||
image_path)
|
||||
logger.info("Detection bbox results save in {}".format(
|
||||
save_name))
|
||||
image.save(save_name, quality=95)
|
||||
|
||||
start = end
|
||||
return results
|
||||
|
||||
def _get_save_image_name(self, output_dir, image_path):
|
||||
"""
|
||||
Get save image name from source image path.
|
||||
"""
|
||||
image_name = os.path.split(image_path)[-1]
|
||||
name, ext = os.path.splitext(image_name)
|
||||
return os.path.join(output_dir, "{}".format(name)) + ext
|
||||
|
||||
def _get_infer_cfg_and_input_spec(self,
|
||||
save_dir,
|
||||
prune_input=True,
|
||||
kl_quant=False):
|
||||
image_shape = None
|
||||
im_shape = [None, 2]
|
||||
scale_factor = [None, 2]
|
||||
test_reader_name = 'TestReader'
|
||||
if 'inputs_def' in self.cfg[test_reader_name]:
|
||||
inputs_def = self.cfg[test_reader_name]['inputs_def']
|
||||
image_shape = inputs_def.get('image_shape', None)
|
||||
# set image_shape=[None, 3, -1, -1] as default
|
||||
if image_shape is None:
|
||||
image_shape = [None, 3, -1, -1]
|
||||
|
||||
if len(image_shape) == 3:
|
||||
image_shape = [None] + image_shape
|
||||
else:
|
||||
im_shape = [image_shape[0], 2]
|
||||
scale_factor = [image_shape[0], 2]
|
||||
|
||||
if hasattr(self.model, 'deploy'):
|
||||
self.model.deploy = True
|
||||
|
||||
for layer in self.model.sublayers():
|
||||
if hasattr(layer, 'convert_to_deploy'):
|
||||
layer.convert_to_deploy()
|
||||
|
||||
if hasattr(self.cfg, 'export') and 'fuse_conv_bn' in self.cfg[
|
||||
'export'] and self.cfg['export']['fuse_conv_bn']:
|
||||
self.model = fuse_conv_bn(self.model)
|
||||
|
||||
export_post_process = self.cfg['export'].get(
|
||||
'post_process', False) if hasattr(self.cfg, 'export') else True
|
||||
export_nms = self.cfg['export'].get('nms', False) if hasattr(
|
||||
self.cfg, 'export') else True
|
||||
export_benchmark = self.cfg['export'].get(
|
||||
'benchmark', False) if hasattr(self.cfg, 'export') else False
|
||||
if hasattr(self.model, 'export_post_process'):
|
||||
self.model.export_post_process = export_post_process if not export_benchmark else False
|
||||
if hasattr(self.model, 'export_nms'):
|
||||
self.model.export_nms = export_nms if not export_benchmark else False
|
||||
if export_post_process and not export_benchmark:
|
||||
image_shape = [None] + image_shape[1:]
|
||||
|
||||
# Save infer cfg
|
||||
_dump_infer_config(self.cfg,
|
||||
os.path.join(save_dir, 'infer_cfg.yml'), image_shape,
|
||||
self.model)
|
||||
|
||||
input_spec = [{
|
||||
"image": InputSpec(
|
||||
shape=image_shape, name='image'),
|
||||
"im_shape": InputSpec(
|
||||
shape=im_shape, name='im_shape'),
|
||||
"scale_factor": InputSpec(
|
||||
shape=scale_factor, name='scale_factor')
|
||||
}]
|
||||
|
||||
if prune_input:
|
||||
static_model = paddle.jit.to_static(
|
||||
self.model, input_spec=input_spec, full_graph=True)
|
||||
# NOTE: dy2st do not pruned program, but jit.save will prune program
|
||||
# input spec, prune input spec here and save with pruned input spec
|
||||
pruned_input_spec = _prune_input_spec(
|
||||
input_spec, static_model.forward.main_program,
|
||||
static_model.forward.outputs)
|
||||
else:
|
||||
static_model = None
|
||||
pruned_input_spec = input_spec
|
||||
|
||||
return static_model, pruned_input_spec
|
||||
|
||||
def export(self, output_dir='output_inference'):
|
||||
if hasattr(self.model, 'aux_neck'):
|
||||
self.model.__delattr__('aux_neck')
|
||||
if hasattr(self.model, 'aux_head'):
|
||||
self.model.__delattr__('aux_head')
|
||||
self.model.eval()
|
||||
|
||||
model_name = os.path.splitext(os.path.split(self.cfg.filename)[-1])[0]
|
||||
save_dir = os.path.join(output_dir, model_name)
|
||||
if not os.path.exists(save_dir):
|
||||
os.makedirs(save_dir)
|
||||
|
||||
static_model, pruned_input_spec = self._get_infer_cfg_and_input_spec(
|
||||
save_dir)
|
||||
|
||||
# dy2st and save model
|
||||
paddle.jit.save(
|
||||
static_model,
|
||||
os.path.join(save_dir, 'model'),
|
||||
input_spec=pruned_input_spec)
|
||||
|
||||
logger.info("Export model and saved in {}".format(save_dir))
|
||||
Reference in New Issue
Block a user