trainer.py 49.4 KB
Newer Older
F
Feng Ni 已提交
1 2 3 4 5 6 7 8 9 10 11 12
# 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
K
Kaipeng Deng 已提交
13 14 15 16 17 18 19
# limitations under the License.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
G
George Ni 已提交
20
import sys
21
import copy
K
Kaipeng Deng 已提交
22
import time
F
Feng Ni 已提交
23
from tqdm import tqdm
M
Manuel Garcia 已提交
24

K
Kaipeng Deng 已提交
25
import numpy as np
M
Mark Ma 已提交
26
import typing
F
Feng Ni 已提交
27
from PIL import Image, ImageOps, ImageFile
W
Wenyu 已提交
28

F
Feng Ni 已提交
29
ImageFile.LOAD_TRUNCATED_IMAGES = True
K
Kaipeng Deng 已提交
30 31

import paddle
F
Feng Ni 已提交
32
import paddle.nn as nn
W
wangguanzhong 已提交
33 34
import paddle.distributed as dist
from paddle.distributed import fleet
K
Kaipeng Deng 已提交
35
from paddle.static import InputSpec
36
from ppdet.optimizer import ModelEMA
K
Kaipeng Deng 已提交
37 38 39

from ppdet.core.workspace import create
from ppdet.utils.checkpoint import load_weight, load_pretrain_weight
C
cnn 已提交
40
from ppdet.utils.visualizer import visualize_results, save_result
41
from ppdet.metrics import Metric, COCOMetric, VOCMetric, WiderFaceMetric, get_infer_results, KeyPointTopDownCOCOEval, KeyPointTopDownMPIIEval, Pose3DEval
42 43
from ppdet.metrics import RBoxMetric, JDEDetMetric, SNIPERCOCOMetric
from ppdet.data.source.sniper_coco import SniperCOCODataSet
K
Kaipeng Deng 已提交
44
from ppdet.data.source.category import get_categories
K
Kaipeng Deng 已提交
45
import ppdet.utils.stats as stats
46
from ppdet.utils.fuse_utils import fuse_conv_bn
47
from ppdet.utils import profiler
48
from ppdet.modeling.post_process import multiclass_nms
K
Kaipeng Deng 已提交
49

50
from .callbacks import Callback, ComposeCallback, LogPrinter, Checkpointer, WiferFaceEval, VisualDLWriter, SniperProposalsGenerator, WandbCallback
G
Guanghua Yu 已提交
51
from .export_utils import _dump_infer_config, _prune_input_spec
K
Kaipeng Deng 已提交
52

53 54
from paddle.distributed.fleet.utils.hybrid_parallel_util import fused_allreduce_gradients

K
Kaipeng Deng 已提交
55
from ppdet.utils.logger import setup_logger
56
logger = setup_logger('ppdet.engine')
K
Kaipeng Deng 已提交
57 58 59

__all__ = ['Trainer']

60
MOT_ARCH = ['DeepSORT', 'JDE', 'FairMOT', 'ByteTrack']
61

K
Kaipeng Deng 已提交
62 63 64 65 66 67 68

class Trainer(object):
    def __init__(self, cfg, mode='train'):
        self.cfg = cfg
        assert mode.lower() in ['train', 'eval', 'test'], \
                "mode should be 'train', 'eval' or 'test'"
        self.mode = mode.lower()
69
        self.optimizer = None
70
        self.is_loaded_weights = False
S
shangliang Xu 已提交
71 72
        self.use_amp = self.cfg.get('amp', False)
        self.amp_level = self.cfg.get('amp_level', 'O1')
73 74
        self.custom_white_list = self.cfg.get('custom_white_list', None)
        self.custom_black_list = self.cfg.get('custom_black_list', None)
K
Kaipeng Deng 已提交
75

G
George Ni 已提交
76
        # build data loader
W
wangguanzhong 已提交
77
        capital_mode = self.mode.capitalize()
78
        if cfg.architecture in MOT_ARCH and self.mode in ['eval', 'test']:
W
wangguanzhong 已提交
79 80
            self.dataset = self.cfg['{}MOTDataset'.format(
                capital_mode)] = create('{}MOTDataset'.format(capital_mode))()
81
        else:
W
wangguanzhong 已提交
82 83
            self.dataset = self.cfg['{}Dataset'.format(capital_mode)] = create(
                '{}Dataset'.format(capital_mode))()
84 85 86 87 88

        if cfg.architecture == 'DeepSORT' and self.mode == 'train':
            logger.error('DeepSORT has no need of training on mot dataset.')
            sys.exit(1)

89 90 91 92
        if cfg.architecture == 'FairMOT' and self.mode == 'eval':
            images = self.parse_mot_images(cfg)
            self.dataset.set_images(images)

G
George Ni 已提交
93
        if self.mode == 'train':
W
wangguanzhong 已提交
94
            self.loader = create('{}Reader'.format(capital_mode))(
G
George Ni 已提交
95 96 97 98
                self.dataset, cfg.worker_num)

        if cfg.architecture == 'JDE' and self.mode == 'train':
            cfg['JDEEmbeddingHead'][
99 100
                'num_identities'] = self.dataset.num_identities_dict[0]
            # JDE only support single class MOT now.
G
George Ni 已提交
101

F
FlyingQianMM 已提交
102
        if cfg.architecture == 'FairMOT' and self.mode == 'train':
M
minghaoBD 已提交
103 104
            cfg['FairMOTEmbeddingHead'][
                'num_identities_dict'] = self.dataset.num_identities_dict
105
            # FairMOT support single class and multi-class MOT now.
F
FlyingQianMM 已提交
106

K
Kaipeng Deng 已提交
107
        # build model
108 109 110 111 112
        if 'model' not in self.cfg:
            self.model = create(cfg.architecture)
        else:
            self.model = self.cfg.model
            self.is_loaded_weights = True
113

F
Feng Ni 已提交
114 115 116
        if cfg.architecture == 'YOLOX':
            for k, m in self.model.named_sublayers():
                if isinstance(m, nn.BatchNorm2D):
F
Feng Ni 已提交
117 118
                    m._epsilon = 1e-3  # for amp(fp16)
                    m._momentum = 0.97  # 0.03 in pytorch
F
Feng Ni 已提交
119

120
        #normalize params for deploy
C
Chang Xu 已提交
121 122 123
        if 'slim' in cfg and cfg['slim_type'] == 'OFA':
            self.model.model.load_meanstd(cfg['TestReader'][
                'sample_transforms'])
C
Chang Xu 已提交
124 125 126 127 128 129 130
        elif 'slim' in cfg and cfg['slim_type'] == 'Distill':
            self.model.student_model.load_meanstd(cfg['TestReader'][
                'sample_transforms'])
        elif 'slim' in cfg and cfg[
                'slim_type'] == 'DistillPrune' and self.mode == 'train':
            self.model.student_model.load_meanstd(cfg['TestReader'][
                'sample_transforms'])
C
Chang Xu 已提交
131 132
        else:
            self.model.load_meanstd(cfg['TestReader']['sample_transforms'])
133

K
Kaipeng Deng 已提交
134 135 136
        # EvalDataset build with BatchSampler to evaluate in single device
        # TODO: multi-device evaluate
        if self.mode == 'eval':
137 138
            if cfg.architecture == 'FairMOT':
                self.loader = create('EvalMOTReader')(self.dataset, 0)
139 140 141
            elif cfg.architecture == "METRO_Body":
                reader_name = '{}Reader'.format(self.mode.capitalize())
                self.loader = create(reader_name)(self.dataset, cfg.worker_num)
142 143 144 145 146 147 148 149 150
            else:
                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':
                    cfg[reader_name]['collate_batch'] = False
                self.loader = create(reader_name)(self.dataset, cfg.worker_num,
                                                  self._eval_batch_sampler)
K
Kaipeng Deng 已提交
151
        # TestDataset build after user set images, skip loader creation here
K
Kaipeng Deng 已提交
152 153 154 155

        # build optimizer in train mode
        if self.mode == 'train':
            steps_per_epoch = len(self.loader)
156 157 158 159
            if steps_per_epoch < 1:
                logger.warning(
                    "Samples in dataset are less than batch_size, please set smaller batch_size in TrainReader."
                )
K
Kaipeng Deng 已提交
160
            self.lr = create('LearningRate')(steps_per_epoch)
W
Wenyu 已提交
161
            self.optimizer = create('OptimizerBuilder')(self.lr, self.model)
K
Kaipeng Deng 已提交
162

M
minghaoBD 已提交
163 164 165 166
            # Unstructured pruner is only enabled in the train mode.
            if self.cfg.get('unstructured_prune'):
                self.pruner = create('UnstructuredPruner')(self.model,
                                                           steps_per_epoch)
S
shangliang Xu 已提交
167
        if self.use_amp and self.amp_level == 'O2':
168 169 170 171
            self.model, self.optimizer = paddle.amp.decorate(
                models=self.model,
                optimizers=self.optimizer,
                level=self.amp_level)
S
shangliang Xu 已提交
172 173 174 175
        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')
176 177
            cycle_epoch = self.cfg.get('cycle_epoch', -1)
            ema_black_list = self.cfg.get('ema_black_list', None)
S
shangliang Xu 已提交
178 179 180 181
            self.ema = ModelEMA(
                self.model,
                decay=ema_decay,
                ema_decay_type=ema_decay_type,
182 183
                cycle_epoch=cycle_epoch,
                ema_black_list=ema_black_list)
S
shangliang Xu 已提交
184

W
wangguanzhong 已提交
185 186
        self._nranks = dist.get_world_size()
        self._local_rank = dist.get_rank()
K
Kaipeng Deng 已提交
187

K
Kaipeng Deng 已提交
188 189 190
        self.status = {}

        self.start_epoch = 0
G
George Ni 已提交
191
        self.end_epoch = 0 if 'epoch' not in cfg else cfg.epoch
K
Kaipeng Deng 已提交
192 193 194 195 196 197 198 199 200 201 202

        # 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)]
203
            if self.cfg.get('use_vdl', False):
204
                self._callbacks.append(VisualDLWriter(self))
205 206
            if self.cfg.get('save_proposals', False):
                self._callbacks.append(SniperProposalsGenerator(self))
207 208
            if self.cfg.get('use_wandb', False) or 'wandb' in self.cfg:
                self._callbacks.append(WandbCallback(self))
K
Kaipeng Deng 已提交
209 210 211
            self._compose_callback = ComposeCallback(self._callbacks)
        elif self.mode == 'eval':
            self._callbacks = [LogPrinter(self)]
212 213
            if self.cfg.metric == 'WiderFace':
                self._callbacks.append(WiferFaceEval(self))
K
Kaipeng Deng 已提交
214
            self._compose_callback = ComposeCallback(self._callbacks)
215
        elif self.mode == 'test' and self.cfg.get('use_vdl', False):
216 217
            self._callbacks = [VisualDLWriter(self)]
            self._compose_callback = ComposeCallback(self._callbacks)
K
Kaipeng Deng 已提交
218 219 220 221
        else:
            self._callbacks = []
            self._compose_callback = None

K
Kaipeng Deng 已提交
222 223
    def _init_metrics(self, validate=False):
        if self.mode == 'test' or (self.mode == 'train' and not validate):
G
Guanghua Yu 已提交
224 225
            self._metrics = []
            return
226
        classwise = self.cfg['classwise'] if 'classwise' in self.cfg else False
227
        if self.cfg.metric == 'COCO' or self.cfg.metric == "SNIPERCOCO":
W
wangxinxin08 已提交
228
            # TODO: bias should be unified
W
wangxinxin08 已提交
229
            bias = 1 if self.cfg.get('bias', False) else 0
S
shangliang Xu 已提交
230 231
            output_eval = self.cfg['output_eval'] \
                if 'output_eval' in self.cfg else None
232
            save_prediction_only = self.cfg.get('save_prediction_only', False)
233 234 235

            # pass clsid2catid info to metric instance to avoid multiple loading
            # annotation file
K
Kaipeng Deng 已提交
236 237
            clsid2catid = {v: k for k, v in self.dataset.catid2clsid.items()} \
                                if self.mode == 'eval' else None
238 239 240 241 242 243 244

            # 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()
245
                dataset = eval_dataset
W
Wenyu 已提交
246 247 248
            else:
                dataset = self.dataset
                anno_file = dataset.get_anno()
249

250
            IouType = self.cfg['IouType'] if 'IouType' in self.cfg else 'bbox'
251 252 253 254 255 256 257 258 259 260 261
            if self.cfg.metric == "COCO":
                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)
                ]
262
            elif self.cfg.metric == "SNIPERCOCO":  # sniper
263 264 265 266 267 268 269 270 271
                self._metrics = [
                    SNIPERCOCOMetric(
                        anno_file=anno_file,
                        dataset=dataset,
                        clsid2catid=clsid2catid,
                        classwise=classwise,
                        output_eval=output_eval,
                        bias=bias,
                        IouType=IouType,
272
                        save_prediction_only=save_prediction_only)
273
                ]
274 275 276 277 278 279
        elif self.cfg.metric == 'RBOX':
            # TODO: bias should be unified
            bias = self.cfg['bias'] if 'bias' in self.cfg 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)
W
wangxinxin08 已提交
280
            imid2path = self.cfg.get('imid2path', None)
281 282 283 284 285 286 287 288 289 290 291 292 293 294 295

            # when do validation in train, annotation file should be get from
            # EvalReader instead of self.dataset(which is TrainReader)
            anno_file = self.dataset.get_anno()
            if self.mode == 'train' and validate:
                eval_dataset = self.cfg['EvalDataset']
                eval_dataset.check_or_download_dataset()
                anno_file = eval_dataset.get_anno()

            self._metrics = [
                RBoxMetric(
                    anno_file=anno_file,
                    classwise=classwise,
                    output_eval=output_eval,
                    bias=bias,
W
wangxinxin08 已提交
296 297
                    save_prediction_only=save_prediction_only,
                    imid2path=imid2path)
298
            ]
K
Kaipeng Deng 已提交
299
        elif self.cfg.metric == 'VOC':
300 301 302 303
            output_eval = self.cfg['output_eval'] \
                if 'output_eval' in self.cfg else None
            save_prediction_only = self.cfg.get('save_prediction_only', False)

K
Kaipeng Deng 已提交
304 305
            self._metrics = [
                VOCMetric(
306
                    label_list=self.dataset.get_label_list(),
K
Kaipeng Deng 已提交
307
                    class_num=self.cfg.num_classes,
308
                    map_type=self.cfg.map_type,
309 310 311
                    classwise=classwise,
                    output_eval=output_eval,
                    save_prediction_only=save_prediction_only)
K
Kaipeng Deng 已提交
312
            ]
313 314 315 316 317 318 319 320 321
        elif self.cfg.metric == 'WiderFace':
            multi_scale = self.cfg.multi_scale_eval if 'multi_scale_eval' in self.cfg else True
            self._metrics = [
                WiderFaceMetric(
                    image_dir=os.path.join(self.dataset.dataset_dir,
                                           self.dataset.image_dir),
                    anno_file=self.dataset.get_anno(),
                    multi_scale=multi_scale)
            ]
322 323 324 325
        elif self.cfg.metric == 'KeyPointTopDownCOCOEval':
            eval_dataset = self.cfg['EvalDataset']
            eval_dataset.check_or_download_dataset()
            anno_file = eval_dataset.get_anno()
326
            save_prediction_only = self.cfg.get('save_prediction_only', False)
327
            self._metrics = [
328 329 330 331 332 333
                KeyPointTopDownCOCOEval(
                    anno_file,
                    len(eval_dataset),
                    self.cfg.num_joints,
                    self.cfg.save_dir,
                    save_prediction_only=save_prediction_only)
334
            ]
Z
zhiboniu 已提交
335 336 337 338
        elif self.cfg.metric == 'KeyPointTopDownMPIIEval':
            eval_dataset = self.cfg['EvalDataset']
            eval_dataset.check_or_download_dataset()
            anno_file = eval_dataset.get_anno()
339
            save_prediction_only = self.cfg.get('save_prediction_only', False)
Z
zhiboniu 已提交
340
            self._metrics = [
341 342 343 344 345 346
                KeyPointTopDownMPIIEval(
                    anno_file,
                    len(eval_dataset),
                    self.cfg.num_joints,
                    self.cfg.save_dir,
                    save_prediction_only=save_prediction_only)
Z
zhiboniu 已提交
347
            ]
348 349 350 351 352 353 354
        elif self.cfg.metric == 'Pose3DEval':
            save_prediction_only = self.cfg.get('save_prediction_only', False)
            self._metrics = [
                Pose3DEval(
                    self.cfg.save_dir,
                    save_prediction_only=save_prediction_only)
            ]
G
George Ni 已提交
355 356
        elif self.cfg.metric == 'MOTDet':
            self._metrics = [JDEDetMetric(), ]
K
Kaipeng Deng 已提交
357
        else:
358
            logger.warning("Metric not support for metric type {}".format(
K
Kaipeng Deng 已提交
359
                self.cfg.metric))
K
Kaipeng Deng 已提交
360 361 362 363 364 365 366
            self._metrics = []

    def _reset_metrics(self):
        for metric in self._metrics:
            metric.reset()

    def register_callbacks(self, callbacks):
367
        callbacks = [c for c in list(callbacks) if c is not None]
K
Kaipeng Deng 已提交
368 369 370 371 372 373 374 375 376 377 378 379 380
        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)

K
Kaipeng Deng 已提交
381
    def load_weights(self, weights):
382 383
        if self.is_loaded_weights:
            return
K
Kaipeng Deng 已提交
384
        self.start_epoch = 0
385
        load_pretrain_weight(self.model, weights)
K
Kaipeng Deng 已提交
386 387
        logger.debug("Load weights {} to start training".format(weights))

388 389 390 391 392 393 394
    def load_weights_sde(self, det_weights, reid_weights):
        if self.model.detector:
            load_weight(self.model.detector, det_weights)
            load_weight(self.model.reid, reid_weights)
        else:
            load_weight(self.model.reid, reid_weights)

K
Kaipeng Deng 已提交
395
    def resume_weights(self, weights):
396 397 398 399 400
        # support Distill resume weights
        if hasattr(self.model, 'student_model'):
            self.start_epoch = load_weight(self.model.student_model, weights,
                                           self.optimizer)
        else:
S
shangliang Xu 已提交
401 402
            self.start_epoch = load_weight(self.model, weights, self.optimizer,
                                           self.ema if self.use_ema else None)
K
Kaipeng Deng 已提交
403
        logger.debug("Resume weights of epoch {}".format(self.start_epoch))
K
Kaipeng Deng 已提交
404

K
Kaipeng Deng 已提交
405
    def train(self, validate=False):
K
Kaipeng Deng 已提交
406
        assert self.mode == 'train', "Model not in 'train' mode"
Z
zhiboniu 已提交
407
        Init_mark = False
W
wangguanzhong 已提交
408
        if validate:
W
wangguanzhong 已提交
409 410
            self.cfg['EvalDataset'] = self.cfg.EvalDataset = create(
                "EvalDataset")()
K
Kaipeng Deng 已提交
411

412
        model = self.model
413
        sync_bn = (getattr(self.cfg, 'norm_type', None) == 'sync_bn' and
414
                   (self.cfg.use_gpu or self.cfg.use_mlu) and self._nranks > 1)
W
wangxinxin08 已提交
415
        if sync_bn:
416
            model = paddle.nn.SyncBatchNorm.convert_sync_batchnorm(model)
W
wangxinxin08 已提交
417

418
        # enabel auto mixed precision mode
S
shangliang Xu 已提交
419
        if self.use_amp:
420
            scaler = paddle.amp.GradScaler(
421
                enable=self.cfg.use_gpu or self.cfg.use_npu or self.cfg.use_mlu,
422 423
                init_loss_scaling=self.cfg.get('init_loss_scaling', 1024))
        # get distributed model
424
        if self.cfg.get('fleet', False):
425
            model = fleet.distributed_model(model)
W
wangguanzhong 已提交
426
            self.optimizer = fleet.distributed_optimizer(self.optimizer)
427
        elif self._nranks > 1:
G
George Ni 已提交
428 429 430
            find_unused_parameters = self.cfg[
                'find_unused_parameters'] if 'find_unused_parameters' in self.cfg else False
            model = paddle.DataParallel(
431
                model, find_unused_parameters=find_unused_parameters)
K
Kaipeng Deng 已提交
432

K
Kaipeng Deng 已提交
433 434 435 436 437 438 439 440 441 442 443 444
        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_staus'] = stats.TrainingStats(self.cfg.log_iter)

G
Guanghua Yu 已提交
445
        if self.cfg.get('print_flops', False):
G
Guanghua Yu 已提交
446 447 448
            flops_loader = create('{}Reader'.format(self.mode.capitalize()))(
                self.dataset, self.cfg.worker_num)
            self._flops(flops_loader)
449
        profiler_options = self.cfg.get('profiler_options', None)
G
Guanghua Yu 已提交
450

451 452
        self._compose_callback.on_train_begin(self.status)

453 454 455
        use_fused_allreduce_gradients = self.cfg[
            'use_fused_allreduce_gradients'] if 'use_fused_allreduce_gradients' in self.cfg else False

K
Kaipeng Deng 已提交
456
        for epoch_id in range(self.start_epoch, self.cfg.epoch):
K
Kaipeng Deng 已提交
457
            self.status['mode'] = 'train'
K
Kaipeng Deng 已提交
458 459 460
            self.status['epoch_id'] = epoch_id
            self._compose_callback.on_epoch_begin(self.status)
            self.loader.dataset.set_epoch(epoch_id)
K
Kaipeng Deng 已提交
461
            model.train()
K
Kaipeng Deng 已提交
462 463 464 465
            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
466
                profiler.add_profiler_step(profiler_options)
K
Kaipeng Deng 已提交
467
                self._compose_callback.on_step_begin(self.status)
S
shangliang Xu 已提交
468
                data['epoch_id'] = epoch_id
K
Kaipeng Deng 已提交
469

S
shangliang Xu 已提交
470
                if self.use_amp:
471 472 473 474
                    if isinstance(
                            model, paddle.
                            DataParallel) and use_fused_allreduce_gradients:
                        with model.no_sync():
F
Feng Ni 已提交
475
                            with paddle.amp.auto_cast(
476
                                    enable=self.cfg.use_gpu or self.cfg.use_mlu,
477 478
                                    custom_white_list=self.custom_white_list,
                                    custom_black_list=self.custom_black_list,
479 480 481 482 483 484 485 486 487 488
                                    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:
F
Feng Ni 已提交
489
                        with paddle.amp.auto_cast(
490
                                enable=self.cfg.use_gpu or self.cfg.use_mlu,
491 492 493
                                custom_white_list=self.custom_white_list,
                                custom_black_list=self.custom_black_list,
                                level=self.amp_level):
494 495 496 497 498 499
                            # model forward
                            outputs = model(data)
                            loss = outputs['loss']
                        # model backward
                        scaled_loss = scaler.scale(loss)
                        scaled_loss.backward()
500 501 502
                    # in dygraph mode, optimizer.minimize is equal to optimizer.step
                    scaler.minimize(self.optimizer, scaled_loss)
                else:
503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519
                    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()
520
                    self.optimizer.step()
K
Kaipeng Deng 已提交
521 522
                curr_lr = self.optimizer.get_lr()
                self.lr.step()
M
minghaoBD 已提交
523 524
                if self.cfg.get('unstructured_prune'):
                    self.pruner.step()
K
Kaipeng Deng 已提交
525 526 527
                self.optimizer.clear_grad()
                self.status['learning_rate'] = curr_lr

K
Kaipeng Deng 已提交
528
                if self._nranks < 2 or self._local_rank == 0:
K
Kaipeng Deng 已提交
529 530 531 532
                    self.status['training_staus'].update(outputs)

                self.status['batch_time'].update(time.time() - iter_tic)
                self._compose_callback.on_step_end(self.status)
533
                if self.use_ema:
S
shangliang Xu 已提交
534
                    self.ema.update()
F
Feng Ni 已提交
535
                iter_tic = time.time()
K
Kaipeng Deng 已提交
536

M
minghaoBD 已提交
537 538
            if self.cfg.get('unstructured_prune'):
                self.pruner.update_params()
539

540
            is_snapshot = (self._nranks < 2 or (self._local_rank == 0 or self.cfg.metric == "Pose3DEval")) \
S
shangliang Xu 已提交
541 542 543 544 545 546 547
                       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

K
Kaipeng Deng 已提交
548 549
            self._compose_callback.on_epoch_end(self.status)

550
            if validate and is_snapshot:
K
Kaipeng Deng 已提交
551 552 553 554 555 556 557
                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'])
558 559 560
                    # If metric is VOC, need to be set collate_batch=False.
                    if self.cfg.metric == 'VOC':
                        self.cfg['EvalReader']['collate_batch'] = False
561 562 563 564 565 566 567 568
                    if self.cfg.metric == "Pose3DEval":
                        self._eval_loader = create('EvalReader')(
                            self._eval_dataset, self.cfg.worker_num)
                    else:
                        self._eval_loader = create('EvalReader')(
                            self._eval_dataset,
                            self.cfg.worker_num,
                            batch_sampler=self._eval_batch_sampler)
Z
zhiboniu 已提交
569 570 571 572 573 574
                # 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()
S
shangliang Xu 已提交
575

K
Kaipeng Deng 已提交
576
                with paddle.no_grad():
577
                    self.status['save_best_model'] = True
K
Kaipeng Deng 已提交
578 579
                    self._eval_with_loader(self._eval_loader)

S
shangliang Xu 已提交
580 581
            if is_snapshot and self.use_ema:
                # reset original weight
582
                self.model.set_dict(weight)
S
shangliang Xu 已提交
583
                self.status.pop('weight')
584

585 586
        self._compose_callback.on_train_end(self.status)

K
Kaipeng Deng 已提交
587
    def _eval_with_loader(self, loader):
K
Kaipeng Deng 已提交
588 589 590
        sample_num = 0
        tic = time.time()
        self._compose_callback.on_epoch_begin(self.status)
K
Kaipeng Deng 已提交
591
        self.status['mode'] = 'eval'
592

K
Kaipeng Deng 已提交
593
        self.model.eval()
G
Guanghua Yu 已提交
594
        if self.cfg.get('print_flops', False):
G
Guanghua Yu 已提交
595 596 597
            flops_loader = create('{}Reader'.format(self.mode.capitalize()))(
                self.dataset, self.cfg.worker_num, self._eval_batch_sampler)
            self._flops(flops_loader)
F
Feng Ni 已提交
598
        for step_id, data in enumerate(loader):
K
Kaipeng Deng 已提交
599 600 601
            self.status['step_id'] = step_id
            self._compose_callback.on_step_begin(self.status)
            # forward
S
shangliang Xu 已提交
602 603
            if self.use_amp:
                with paddle.amp.auto_cast(
604
                        enable=self.cfg.use_gpu or self.cfg.use_mlu,
605 606 607
                        custom_white_list=self.custom_white_list,
                        custom_black_list=self.custom_black_list,
                        level=self.amp_level):
S
shangliang Xu 已提交
608 609 610
                    outs = self.model(data)
            else:
                outs = self.model(data)
K
Kaipeng Deng 已提交
611 612 613 614 615

            # update metrics
            for metric in self._metrics:
                metric.update(data, outs)

M
Mark Ma 已提交
616 617 618 619 620
            # 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]
K
Kaipeng Deng 已提交
621 622 623 624 625 626 627 628 629
            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()
630
        self._compose_callback.on_epoch_end(self.status)
K
Kaipeng Deng 已提交
631 632 633
        # reset metric states for metric may performed multiple times
        self._reset_metrics()

K
Kaipeng Deng 已提交
634
    def evaluate(self):
635 636 637 638 639 640 641 642 643
        # 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)
644 645
        with paddle.no_grad():
            self._eval_with_loader(self.loader)
K
Kaipeng Deng 已提交
646

647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670
    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()
        if self.cfg.get('print_flops', False):
            flops_loader = create('{}Reader'.format(self.mode.capitalize()))(
                self.dataset, self.cfg.worker_num, self._eval_batch_sampler)
            self._flops(flops_loader)

        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(
671
                        enable=self.cfg.use_gpu or self.cfg.use_mlu,
672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697 698 699 700 701 702 703 704 705 706 707 708 709 710 711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746
                        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',
F
Feng Ni 已提交
747 748
                      save_results=False,
                      visualize=True):
749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 767 768 769 770 771 772 773 774 775 776 777 778 779 780 781 782 783 784 785 786 787 788 789 790 791 792 793 794 795 796 797 798 799 800
        self.dataset.set_slice_images(images, slice_size, overlap_ratio)
        loader = create('TestReader')(self.dataset, 0)

        imid2path = self.dataset.get_imid2path()

        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()
        if self.cfg.get('print_flops', False):
            flops_loader = create('TestReader')(self.dataset, 0)
            self._flops(flops_loader)

        results = []  # all images
        merged_bboxs = []  # single image
        for step_id, data in enumerate(tqdm(loader)):
            self.status['step_id'] = step_id
            # forward
            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 key in ['im_shape', 'scale_factor', 'im_id']:
                    if isinstance(data, typing.Sequence):
F
Feng Ni 已提交
801
                        merged_results[key] = data[0][key]
802
                    else:
F
Feng Ni 已提交
803
                        merged_results[key] = data[key]
804 805 806 807 808
                for key, value in merged_results.items():
                    if hasattr(value, 'numpy'):
                        merged_results[key] = value.numpy()
                results.append(merged_results)

F
Feng Ni 已提交
809 810 811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834 835
        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, segm_res, keypoint_res = None, None, None
                    image = visualize_results(
                        image, bbox_res, mask_res, segm_res, keypoint_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
836

C
cnn 已提交
837 838 839 840
    def predict(self,
                images,
                draw_threshold=0.5,
                output_dir='output',
W
wangxinxin08 已提交
841 842 843 844 845
                save_results=False,
                visualize=True):
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

K
Kaipeng Deng 已提交
846 847 848
        self.dataset.set_images(images)
        loader = create('TestReader')(self.dataset, 0)

W
wangxinxin08 已提交
849 850
        imid2path = self.dataset.get_imid2path()

W
Wenyu 已提交
851 852 853 854 855 856 857 858 859 860 861 862 863
        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
W
wangxinxin08 已提交
864
            self.cfg['imid2path'] = imid2path
W
Wenyu 已提交
865 866 867 868 869 870 871 872 873 874 875 876
            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

W
wangxinxin08 已提交
877 878
            self.cfg.pop('imid2path')

W
Wenyu 已提交
879 880 881 882 883 884 885 886 887 888
            _metrics = copy.deepcopy(self._metrics)
            self._metrics = metrics

            return _metrics

        if save_results:
            metrics = setup_metrics_for_loader()
        else:
            metrics = []

K
Kaipeng Deng 已提交
889
        anno_file = self.dataset.get_anno()
C
cnn 已提交
890 891
        clsid2catid, catid2name = get_categories(
            self.cfg.metric, anno_file=anno_file)
K
Kaipeng Deng 已提交
892

K
Kaipeng Deng 已提交
893 894 895
        # Run Infer 
        self.status['mode'] = 'test'
        self.model.eval()
G
Guanghua Yu 已提交
896
        if self.cfg.get('print_flops', False):
G
Guanghua Yu 已提交
897 898
            flops_loader = create('TestReader')(self.dataset, 0)
            self._flops(flops_loader)
899
        results = []
F
Feng Ni 已提交
900
        for step_id, data in enumerate(tqdm(loader)):
K
Kaipeng Deng 已提交
901 902 903
            self.status['step_id'] = step_id
            # forward
            outs = self.model(data)
904

W
Wenyu 已提交
905 906 907
            for _m in metrics:
                _m.update(data, outs)

K
Kaipeng Deng 已提交
908
            for key in ['im_shape', 'scale_factor', 'im_id']:
M
Mark Ma 已提交
909 910 911 912
                if isinstance(data, typing.Sequence):
                    outs[key] = data[0][key]
                else:
                    outs[key] = data[key]
G
Guanghua Yu 已提交
913
            for key, value in outs.items():
914 915
                if hasattr(value, 'numpy'):
                    outs[key] = value.numpy()
916
            results.append(outs)
W
Wenyu 已提交
917

918 919
        # sniper
        if type(self.dataset) == SniperCOCODataSet:
920 921
            results = self.dataset.anno_cropper.aggregate_chips_detections(
                results)
K
Kaipeng Deng 已提交
922

W
Wenyu 已提交
923 924 925 926
        for _m in metrics:
            _m.accumulate()
            _m.reset()

W
wangxinxin08 已提交
927 928 929 930 931 932 933 934 935 936 937 938 939 940 941 942 943 944 945 946 947
        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
948 949
                    pose3d_res = batch_res['pose3d'][start:end] \
                            if 'pose3d' in batch_res else None
W
wangxinxin08 已提交
950 951
                    image = visualize_results(
                        image, bbox_res, mask_res, segm_res, keypoint_res,
952
                        pose3d_res, int(im_id), catid2name, draw_threshold)
W
wangxinxin08 已提交
953 954 955 956 957 958 959 960 961 962 963
                    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
K
Kaipeng Deng 已提交
964 965 966 967 968 969 970 971 972

    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

S
shangliang Xu 已提交
973 974 975 976
    def _get_infer_cfg_and_input_spec(self,
                                      save_dir,
                                      prune_input=True,
                                      kl_quant=False):
K
Kaipeng Deng 已提交
977
        image_shape = None
978 979
        im_shape = [None, 2]
        scale_factor = [None, 2]
980 981 982 983 984 985
        if self.cfg.architecture in MOT_ARCH:
            test_reader_name = 'TestMOTReader'
        else:
            test_reader_name = 'TestReader'
        if 'inputs_def' in self.cfg[test_reader_name]:
            inputs_def = self.cfg[test_reader_name]['inputs_def']
K
Kaipeng Deng 已提交
986
            image_shape = inputs_def.get('image_shape', None)
G
Guanghua Yu 已提交
987
        # set image_shape=[None, 3, -1, -1] as default
K
Kaipeng Deng 已提交
988
        if image_shape is None:
G
Guanghua Yu 已提交
989
            image_shape = [None, 3, -1, -1]
990

G
Guanghua Yu 已提交
991 992
        if len(image_shape) == 3:
            image_shape = [None] + image_shape
993 994 995
        else:
            im_shape = [image_shape[0], 2]
            scale_factor = [image_shape[0], 2]
K
Kaipeng Deng 已提交
996

997
        if hasattr(self.model, 'deploy'):
998
            self.model.deploy = True
S
shangliang Xu 已提交
999

1000 1001 1002 1003
        if 'slim' not in self.cfg:
            for layer in self.model.sublayers():
                if hasattr(layer, 'convert_to_deploy'):
                    layer.convert_to_deploy()
S
shangliang Xu 已提交
1004

1005 1006 1007 1008 1009 1010
        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
1011 1012 1013
        if hasattr(self.model, 'fuse_norm'):
            self.model.fuse_norm = self.cfg['TestReader'].get('fuse_normalize',
                                                              False)
1014 1015 1016 1017 1018 1019
        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:]
K
Kaipeng Deng 已提交
1020

K
Kaipeng Deng 已提交
1021 1022 1023 1024 1025 1026 1027
        # Save infer cfg
        _dump_infer_config(self.cfg,
                           os.path.join(save_dir, 'infer_cfg.yml'), image_shape,
                           self.model)

        input_spec = [{
            "image": InputSpec(
G
Guanghua Yu 已提交
1028
                shape=image_shape, name='image'),
K
Kaipeng Deng 已提交
1029
            "im_shape": InputSpec(
1030
                shape=im_shape, name='im_shape'),
K
Kaipeng Deng 已提交
1031
            "scale_factor": InputSpec(
1032
                shape=scale_factor, name='scale_factor')
K
Kaipeng Deng 已提交
1033
        }]
G
George Ni 已提交
1034 1035 1036 1037 1038
        if self.cfg.architecture == 'DeepSORT':
            input_spec[0].update({
                "crops": InputSpec(
                    shape=[None, 3, 192, 64], name='crops')
            })
G
Guanghua Yu 已提交
1039 1040 1041 1042 1043 1044 1045 1046 1047 1048 1049 1050
        if prune_input:
            static_model = paddle.jit.to_static(
                self.model, input_spec=input_spec)
            # 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

G
Guanghua Yu 已提交
1051
        # TODO: Hard code, delete it when support prune input_spec.
1052
        if self.cfg.architecture == 'PicoDet' and not export_post_process:
G
Guanghua Yu 已提交
1053 1054 1055 1056
            pruned_input_spec = [{
                "image": InputSpec(
                    shape=image_shape, name='image')
            }]
S
shangliang Xu 已提交
1057 1058 1059 1060 1061 1062 1063 1064 1065 1066 1067 1068 1069
        if kl_quant:
            if self.cfg.architecture == 'PicoDet' or 'ppyoloe' in self.cfg.weights:
                pruned_input_spec = [{
                    "image": InputSpec(
                        shape=image_shape, name='image'),
                    "scale_factor": InputSpec(
                        shape=scale_factor, name='scale_factor')
                }]
            elif 'tinypose' in self.cfg.weights:
                pruned_input_spec = [{
                    "image": InputSpec(
                        shape=image_shape, name='image')
                }]
G
Guanghua Yu 已提交
1070

G
Guanghua Yu 已提交
1071 1072 1073 1074
        return static_model, pruned_input_spec

    def export(self, output_dir='output_inference'):
        self.model.eval()
1075 1076 1077 1078 1079

        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)

G
Guanghua Yu 已提交
1080 1081 1082 1083
        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)
K
Kaipeng Deng 已提交
1084

G
Guanghua Yu 已提交
1085 1086
        static_model, pruned_input_spec = self._get_infer_cfg_and_input_spec(
            save_dir)
G
Guanghua Yu 已提交
1087 1088

        # dy2st and save model
1089
        if 'slim' not in self.cfg or 'QAT' not in self.cfg['slim_type']:
1090 1091 1092 1093 1094
            paddle.jit.save(
                static_model,
                os.path.join(save_dir, 'model'),
                input_spec=pruned_input_spec)
        else:
1095
            self.cfg.slim.save_quantized_model(
1096 1097
                self.model,
                os.path.join(save_dir, 'model'),
G
Guanghua Yu 已提交
1098 1099
                input_spec=pruned_input_spec)
        logger.info("Export model and saved in {}".format(save_dir))
1100

G
Guanghua Yu 已提交
1101 1102 1103 1104 1105 1106 1107 1108 1109 1110 1111 1112
    def post_quant(self, output_dir='output_inference'):
        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)

        for idx, data in enumerate(self.loader):
            self.model(data)
            if idx == int(self.cfg.get('quant_batch_num', 10)):
                break

        # TODO: support prune input_spec
S
shangliang Xu 已提交
1113
        kl_quant = True if hasattr(self.cfg.slim, 'ptq') else False
G
Guanghua Yu 已提交
1114
        _, pruned_input_spec = self._get_infer_cfg_and_input_spec(
S
shangliang Xu 已提交
1115
            save_dir, prune_input=False, kl_quant=kl_quant)
G
Guanghua Yu 已提交
1116 1117 1118 1119 1120 1121

        self.cfg.slim.save_quantized_model(
            self.model,
            os.path.join(save_dir, 'model'),
            input_spec=pruned_input_spec)
        logger.info("Export Post-Quant model and saved in {}".format(save_dir))
G
Guanghua Yu 已提交
1122 1123 1124 1125 1126 1127 1128 1129 1130 1131 1132 1133 1134 1135 1136 1137 1138 1139 1140 1141 1142 1143 1144 1145 1146

    def _flops(self, loader):
        self.model.eval()
        try:
            import paddleslim
        except Exception as e:
            logger.warning(
                'Unable to calculate flops, please install paddleslim, for example: `pip install paddleslim`'
            )
            return

        from paddleslim.analysis import dygraph_flops as flops
        input_data = None
        for data in loader:
            input_data = data
            break

        input_spec = [{
            "image": input_data['image'][0].unsqueeze(0),
            "im_shape": input_data['im_shape'][0].unsqueeze(0),
            "scale_factor": input_data['scale_factor'][0].unsqueeze(0)
        }]
        flops = flops(self.model, input_spec) / (1000**3)
        logger.info(" Model FLOPs : {:.6f}G. (image shape is {})".format(
            flops, input_data['image'][0].unsqueeze(0).shape))
1147 1148 1149 1150 1151 1152 1153 1154 1155 1156 1157 1158 1159 1160 1161 1162 1163 1164 1165 1166 1167 1168 1169

    def parse_mot_images(self, cfg):
        import glob
        # for quant
        dataset_dir = cfg['EvalMOTDataset'].dataset_dir
        data_root = cfg['EvalMOTDataset'].data_root
        data_root = '{}/{}'.format(dataset_dir, data_root)
        seqs = os.listdir(data_root)
        seqs.sort()
        all_images = []
        for seq in seqs:
            infer_dir = os.path.join(data_root, seq)
            assert infer_dir is None or os.path.isdir(infer_dir), \
                "{} is not a directory".format(infer_dir)
            images = set()
            exts = ['jpg', 'jpeg', 'png', 'bmp']
            exts += [ext.upper() for ext in exts]
            for ext in exts:
                images.update(glob.glob('{}/*.{}'.format(infer_dir, ext)))
            images = list(images)
            images.sort()
            assert len(images) > 0, "no image found in {}".format(infer_dir)
            all_images.extend(images)
1170 1171 1172
            logger.info("Found {} inference images in total.".format(
                len(images)))
        return all_images