trainer.py 47.8 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
Z
zhiboniu 已提交
41
from ppdet.metrics import Metric, COCOMetric, VOCMetric, WiderFaceMetric, get_infer_results, KeyPointTopDownCOCOEval, KeyPointTopDownMPIIEval
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 139 140 141 142 143 144 145 146 147
            if cfg.architecture == 'FairMOT':
                self.loader = create('EvalMOTReader')(self.dataset, 0)
            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 已提交
148
        # TestDataset build after user set images, skip loader creation here
K
Kaipeng Deng 已提交
149 150 151 152

        # build optimizer in train mode
        if self.mode == 'train':
            steps_per_epoch = len(self.loader)
153 154 155 156
            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 已提交
157
            self.lr = create('LearningRate')(steps_per_epoch)
W
Wenyu 已提交
158
            self.optimizer = create('OptimizerBuilder')(self.lr, self.model)
K
Kaipeng Deng 已提交
159

M
minghaoBD 已提交
160 161 162 163
            # 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 已提交
164
        if self.use_amp and self.amp_level == 'O2':
165 166 167 168
            self.model, self.optimizer = paddle.amp.decorate(
                models=self.model,
                optimizers=self.optimizer,
                level=self.amp_level)
S
shangliang Xu 已提交
169 170 171 172 173 174 175 176 177 178 179
        self.use_ema = ('use_ema' in cfg and cfg['use_ema'])
        if self.use_ema:
            ema_decay = self.cfg.get('ema_decay', 0.9998)
            cycle_epoch = self.cfg.get('cycle_epoch', -1)
            ema_decay_type = self.cfg.get('ema_decay_type', 'threshold')
            self.ema = ModelEMA(
                self.model,
                decay=ema_decay,
                ema_decay_type=ema_decay_type,
                cycle_epoch=cycle_epoch)

W
wangguanzhong 已提交
180 181
        self._nranks = dist.get_world_size()
        self._local_rank = dist.get_rank()
K
Kaipeng Deng 已提交
182

K
Kaipeng Deng 已提交
183 184 185
        self.status = {}

        self.start_epoch = 0
G
George Ni 已提交
186
        self.end_epoch = 0 if 'epoch' not in cfg else cfg.epoch
K
Kaipeng Deng 已提交
187 188 189 190 191 192 193 194 195 196 197

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

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

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

            # 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()
240
                dataset = eval_dataset
W
Wenyu 已提交
241 242 243
            else:
                dataset = self.dataset
                anno_file = dataset.get_anno()
244

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

            # 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 已提交
291 292
                    save_prediction_only=save_prediction_only,
                    imid2path=imid2path)
293
            ]
K
Kaipeng Deng 已提交
294
        elif self.cfg.metric == 'VOC':
295 296 297 298
            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 已提交
299 300
            self._metrics = [
                VOCMetric(
301
                    label_list=self.dataset.get_label_list(),
K
Kaipeng Deng 已提交
302
                    class_num=self.cfg.num_classes,
303
                    map_type=self.cfg.map_type,
304 305 306
                    classwise=classwise,
                    output_eval=output_eval,
                    save_prediction_only=save_prediction_only)
K
Kaipeng Deng 已提交
307
            ]
308 309 310 311 312 313 314 315 316
        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)
            ]
317 318 319 320
        elif self.cfg.metric == 'KeyPointTopDownCOCOEval':
            eval_dataset = self.cfg['EvalDataset']
            eval_dataset.check_or_download_dataset()
            anno_file = eval_dataset.get_anno()
321
            save_prediction_only = self.cfg.get('save_prediction_only', False)
322
            self._metrics = [
323 324 325 326 327 328
                KeyPointTopDownCOCOEval(
                    anno_file,
                    len(eval_dataset),
                    self.cfg.num_joints,
                    self.cfg.save_dir,
                    save_prediction_only=save_prediction_only)
329
            ]
Z
zhiboniu 已提交
330 331 332 333
        elif self.cfg.metric == 'KeyPointTopDownMPIIEval':
            eval_dataset = self.cfg['EvalDataset']
            eval_dataset.check_or_download_dataset()
            anno_file = eval_dataset.get_anno()
334
            save_prediction_only = self.cfg.get('save_prediction_only', False)
Z
zhiboniu 已提交
335
            self._metrics = [
336 337 338 339 340 341
                KeyPointTopDownMPIIEval(
                    anno_file,
                    len(eval_dataset),
                    self.cfg.num_joints,
                    self.cfg.save_dir,
                    save_prediction_only=save_prediction_only)
Z
zhiboniu 已提交
342
            ]
G
George Ni 已提交
343 344
        elif self.cfg.metric == 'MOTDet':
            self._metrics = [JDEDetMetric(), ]
K
Kaipeng Deng 已提交
345
        else:
346
            logger.warning("Metric not support for metric type {}".format(
K
Kaipeng Deng 已提交
347
                self.cfg.metric))
K
Kaipeng Deng 已提交
348 349 350 351 352 353 354
            self._metrics = []

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

    def register_callbacks(self, callbacks):
355
        callbacks = [c for c in list(callbacks) if c is not None]
K
Kaipeng Deng 已提交
356 357 358 359 360 361 362 363 364 365 366 367 368
        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 已提交
369
    def load_weights(self, weights):
370 371
        if self.is_loaded_weights:
            return
K
Kaipeng Deng 已提交
372
        self.start_epoch = 0
373
        load_pretrain_weight(self.model, weights)
K
Kaipeng Deng 已提交
374 375
        logger.debug("Load weights {} to start training".format(weights))

376 377 378 379 380 381 382
    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 已提交
383
    def resume_weights(self, weights):
384 385 386 387 388
        # 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 已提交
389 390
            self.start_epoch = load_weight(self.model, weights, self.optimizer,
                                           self.ema if self.use_ema else None)
K
Kaipeng Deng 已提交
391
        logger.debug("Resume weights of epoch {}".format(self.start_epoch))
K
Kaipeng Deng 已提交
392

K
Kaipeng Deng 已提交
393
    def train(self, validate=False):
K
Kaipeng Deng 已提交
394
        assert self.mode == 'train', "Model not in 'train' mode"
Z
zhiboniu 已提交
395
        Init_mark = False
W
wangguanzhong 已提交
396
        if validate:
W
wangguanzhong 已提交
397 398
            self.cfg['EvalDataset'] = self.cfg.EvalDataset = create(
                "EvalDataset")()
K
Kaipeng Deng 已提交
399

400
        model = self.model
401
        sync_bn = (getattr(self.cfg, 'norm_type', None) == 'sync_bn' and
W
wangxinxin08 已提交
402 403
                   self.cfg.use_gpu and self._nranks > 1)
        if sync_bn:
404
            model = paddle.nn.SyncBatchNorm.convert_sync_batchnorm(model)
W
wangxinxin08 已提交
405

406
        # enabel auto mixed precision mode
S
shangliang Xu 已提交
407
        if self.use_amp:
408 409 410 411
            scaler = paddle.amp.GradScaler(
                enable=self.cfg.use_gpu or self.cfg.use_npu,
                init_loss_scaling=self.cfg.get('init_loss_scaling', 1024))
        # get distributed model
412
        if self.cfg.get('fleet', False):
413
            model = fleet.distributed_model(model)
W
wangguanzhong 已提交
414
            self.optimizer = fleet.distributed_optimizer(self.optimizer)
415
        elif self._nranks > 1:
G
George Ni 已提交
416 417 418
            find_unused_parameters = self.cfg[
                'find_unused_parameters'] if 'find_unused_parameters' in self.cfg else False
            model = paddle.DataParallel(
419
                model, find_unused_parameters=find_unused_parameters)
K
Kaipeng Deng 已提交
420

K
Kaipeng Deng 已提交
421 422 423 424 425 426 427 428 429 430 431 432
        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 已提交
433
        if self.cfg.get('print_flops', False):
G
Guanghua Yu 已提交
434 435 436
            flops_loader = create('{}Reader'.format(self.mode.capitalize()))(
                self.dataset, self.cfg.worker_num)
            self._flops(flops_loader)
437
        profiler_options = self.cfg.get('profiler_options', None)
G
Guanghua Yu 已提交
438

439 440
        self._compose_callback.on_train_begin(self.status)

441 442 443
        use_fused_allreduce_gradients = self.cfg[
            'use_fused_allreduce_gradients'] if 'use_fused_allreduce_gradients' in self.cfg else False

K
Kaipeng Deng 已提交
444
        for epoch_id in range(self.start_epoch, self.cfg.epoch):
K
Kaipeng Deng 已提交
445
            self.status['mode'] = 'train'
K
Kaipeng Deng 已提交
446 447 448
            self.status['epoch_id'] = epoch_id
            self._compose_callback.on_epoch_begin(self.status)
            self.loader.dataset.set_epoch(epoch_id)
K
Kaipeng Deng 已提交
449
            model.train()
K
Kaipeng Deng 已提交
450 451 452 453
            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
454
                profiler.add_profiler_step(profiler_options)
K
Kaipeng Deng 已提交
455
                self._compose_callback.on_step_begin(self.status)
S
shangliang Xu 已提交
456
                data['epoch_id'] = epoch_id
K
Kaipeng Deng 已提交
457

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

K
Kaipeng Deng 已提交
516
                if self._nranks < 2 or self._local_rank == 0:
K
Kaipeng Deng 已提交
517 518 519 520
                    self.status['training_staus'].update(outputs)

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

M
minghaoBD 已提交
525 526
            if self.cfg.get('unstructured_prune'):
                self.pruner.update_params()
527

S
shangliang Xu 已提交
528 529 530 531 532 533 534 535
            is_snapshot = (self._nranks < 2 or self._local_rank == 0) \
                       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 已提交
536 537
            self._compose_callback.on_epoch_end(self.status)

S
shangliang Xu 已提交
538
            if validate and is_snapshot:
K
Kaipeng Deng 已提交
539 540 541 542 543 544 545
                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'])
546 547 548
                    # If metric is VOC, need to be set collate_batch=False.
                    if self.cfg.metric == 'VOC':
                        self.cfg['EvalReader']['collate_batch'] = False
K
Kaipeng Deng 已提交
549 550 551 552
                    self._eval_loader = create('EvalReader')(
                        self._eval_dataset,
                        self.cfg.worker_num,
                        batch_sampler=self._eval_batch_sampler)
Z
zhiboniu 已提交
553 554 555 556 557 558
                # 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 已提交
559

K
Kaipeng Deng 已提交
560
                with paddle.no_grad():
561
                    self.status['save_best_model'] = True
K
Kaipeng Deng 已提交
562 563
                    self._eval_with_loader(self._eval_loader)

S
shangliang Xu 已提交
564 565
            if is_snapshot and self.use_ema:
                # reset original weight
566
                self.model.set_dict(weight)
S
shangliang Xu 已提交
567
                self.status.pop('weight')
568

569 570
        self._compose_callback.on_train_end(self.status)

K
Kaipeng Deng 已提交
571
    def _eval_with_loader(self, loader):
K
Kaipeng Deng 已提交
572 573 574
        sample_num = 0
        tic = time.time()
        self._compose_callback.on_epoch_begin(self.status)
K
Kaipeng Deng 已提交
575 576
        self.status['mode'] = 'eval'
        self.model.eval()
G
Guanghua Yu 已提交
577
        if self.cfg.get('print_flops', False):
G
Guanghua Yu 已提交
578 579 580
            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 已提交
581
        for step_id, data in enumerate(loader):
K
Kaipeng Deng 已提交
582 583 584
            self.status['step_id'] = step_id
            self._compose_callback.on_step_begin(self.status)
            # forward
S
shangliang Xu 已提交
585 586
            if self.use_amp:
                with paddle.amp.auto_cast(
587 588 589 590
                        enable=self.cfg.use_gpu,
                        custom_white_list=self.custom_white_list,
                        custom_black_list=self.custom_black_list,
                        level=self.amp_level):
S
shangliang Xu 已提交
591 592 593
                    outs = self.model(data)
            else:
                outs = self.model(data)
K
Kaipeng Deng 已提交
594 595 596 597 598

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

M
Mark Ma 已提交
599 600 601 602 603
            # 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 已提交
604 605 606 607 608 609 610 611 612
            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()
613
        self._compose_callback.on_epoch_end(self.status)
K
Kaipeng Deng 已提交
614 615 616
        # reset metric states for metric may performed multiple times
        self._reset_metrics()

K
Kaipeng Deng 已提交
617
    def evaluate(self):
618 619
        with paddle.no_grad():
            self._eval_with_loader(self.loader)
K
Kaipeng Deng 已提交
620

621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 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 671 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
    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(
                        enable=self.cfg.use_gpu,
                        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 已提交
721 722
                      save_results=False,
                      visualize=True):
723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 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
        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 已提交
775
                        merged_results[key] = data[0][key]
776
                    else:
F
Feng Ni 已提交
777
                        merged_results[key] = data[key]
778 779 780 781 782
                for key, value in merged_results.items():
                    if hasattr(value, 'numpy'):
                        merged_results[key] = value.numpy()
                results.append(merged_results)

F
Feng Ni 已提交
783 784 785 786 787 788 789 790 791 792 793 794 795 796 797 798 799 800 801 802 803 804 805 806 807 808 809
        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
810

C
cnn 已提交
811 812 813 814
    def predict(self,
                images,
                draw_threshold=0.5,
                output_dir='output',
W
wangxinxin08 已提交
815 816 817 818 819
                save_results=False,
                visualize=True):
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

K
Kaipeng Deng 已提交
820 821 822
        self.dataset.set_images(images)
        loader = create('TestReader')(self.dataset, 0)

W
wangxinxin08 已提交
823 824
        imid2path = self.dataset.get_imid2path()

W
Wenyu 已提交
825 826 827 828 829 830 831 832 833 834 835 836 837
        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 已提交
838
            self.cfg['imid2path'] = imid2path
W
Wenyu 已提交
839 840 841 842 843 844 845 846 847 848 849 850
            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 已提交
851 852
            self.cfg.pop('imid2path')

W
Wenyu 已提交
853 854 855 856 857 858 859 860 861 862
            _metrics = copy.deepcopy(self._metrics)
            self._metrics = metrics

            return _metrics

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

K
Kaipeng Deng 已提交
863
        anno_file = self.dataset.get_anno()
C
cnn 已提交
864 865
        clsid2catid, catid2name = get_categories(
            self.cfg.metric, anno_file=anno_file)
K
Kaipeng Deng 已提交
866

K
Kaipeng Deng 已提交
867 868 869
        # Run Infer 
        self.status['mode'] = 'test'
        self.model.eval()
G
Guanghua Yu 已提交
870
        if self.cfg.get('print_flops', False):
G
Guanghua Yu 已提交
871 872
            flops_loader = create('TestReader')(self.dataset, 0)
            self._flops(flops_loader)
873
        results = []
F
Feng Ni 已提交
874
        for step_id, data in enumerate(tqdm(loader)):
K
Kaipeng Deng 已提交
875 876 877
            self.status['step_id'] = step_id
            # forward
            outs = self.model(data)
878

W
Wenyu 已提交
879 880 881
            for _m in metrics:
                _m.update(data, outs)

K
Kaipeng Deng 已提交
882
            for key in ['im_shape', 'scale_factor', 'im_id']:
M
Mark Ma 已提交
883 884 885 886
                if isinstance(data, typing.Sequence):
                    outs[key] = data[0][key]
                else:
                    outs[key] = data[key]
G
Guanghua Yu 已提交
887
            for key, value in outs.items():
888 889
                if hasattr(value, 'numpy'):
                    outs[key] = value.numpy()
890
            results.append(outs)
W
Wenyu 已提交
891

892 893
        # sniper
        if type(self.dataset) == SniperCOCODataSet:
894 895
            results = self.dataset.anno_cropper.aggregate_chips_detections(
                results)
K
Kaipeng Deng 已提交
896

W
Wenyu 已提交
897 898 899 900
        for _m in metrics:
            _m.accumulate()
            _m.reset()

W
wangxinxin08 已提交
901 902 903 904 905 906 907 908 909 910 911 912 913 914 915 916 917 918 919 920 921 922 923 924 925 926 927 928 929 930 931 932 933 934 935
        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
                    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
K
Kaipeng Deng 已提交
936 937 938 939 940 941 942 943 944

    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 已提交
945 946 947 948
    def _get_infer_cfg_and_input_spec(self,
                                      save_dir,
                                      prune_input=True,
                                      kl_quant=False):
K
Kaipeng Deng 已提交
949
        image_shape = None
950 951
        im_shape = [None, 2]
        scale_factor = [None, 2]
952 953 954 955 956 957
        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 已提交
958
            image_shape = inputs_def.get('image_shape', None)
G
Guanghua Yu 已提交
959
        # set image_shape=[None, 3, -1, -1] as default
K
Kaipeng Deng 已提交
960
        if image_shape is None:
G
Guanghua Yu 已提交
961
            image_shape = [None, 3, -1, -1]
962

G
Guanghua Yu 已提交
963 964
        if len(image_shape) == 3:
            image_shape = [None] + image_shape
965 966 967
        else:
            im_shape = [image_shape[0], 2]
            scale_factor = [image_shape[0], 2]
K
Kaipeng Deng 已提交
968

969
        if hasattr(self.model, 'deploy'):
970
            self.model.deploy = True
S
shangliang Xu 已提交
971

972 973 974 975
        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 已提交
976

977 978 979 980 981 982
        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
983 984 985
        if hasattr(self.model, 'fuse_norm'):
            self.model.fuse_norm = self.cfg['TestReader'].get('fuse_normalize',
                                                              False)
986 987 988 989 990 991
        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 已提交
992

K
Kaipeng Deng 已提交
993 994 995 996 997 998 999
        # 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 已提交
1000
                shape=image_shape, name='image'),
K
Kaipeng Deng 已提交
1001
            "im_shape": InputSpec(
1002
                shape=im_shape, name='im_shape'),
K
Kaipeng Deng 已提交
1003
            "scale_factor": InputSpec(
1004
                shape=scale_factor, name='scale_factor')
K
Kaipeng Deng 已提交
1005
        }]
G
George Ni 已提交
1006 1007 1008 1009 1010
        if self.cfg.architecture == 'DeepSORT':
            input_spec[0].update({
                "crops": InputSpec(
                    shape=[None, 3, 192, 64], name='crops')
            })
G
Guanghua Yu 已提交
1011 1012 1013 1014 1015 1016 1017 1018 1019 1020 1021 1022
        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 已提交
1023
        # TODO: Hard code, delete it when support prune input_spec.
1024
        if self.cfg.architecture == 'PicoDet' and not export_post_process:
G
Guanghua Yu 已提交
1025 1026 1027 1028
            pruned_input_spec = [{
                "image": InputSpec(
                    shape=image_shape, name='image')
            }]
S
shangliang Xu 已提交
1029 1030 1031 1032 1033 1034 1035 1036 1037 1038 1039 1040 1041
        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 已提交
1042

G
Guanghua Yu 已提交
1043 1044 1045 1046
        return static_model, pruned_input_spec

    def export(self, output_dir='output_inference'):
        self.model.eval()
1047 1048 1049 1050 1051

        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 已提交
1052 1053 1054 1055
        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 已提交
1056

G
Guanghua Yu 已提交
1057 1058
        static_model, pruned_input_spec = self._get_infer_cfg_and_input_spec(
            save_dir)
G
Guanghua Yu 已提交
1059 1060

        # dy2st and save model
1061
        if 'slim' not in self.cfg or 'QAT' not in self.cfg['slim_type']:
1062 1063 1064 1065 1066
            paddle.jit.save(
                static_model,
                os.path.join(save_dir, 'model'),
                input_spec=pruned_input_spec)
        else:
1067
            self.cfg.slim.save_quantized_model(
1068 1069
                self.model,
                os.path.join(save_dir, 'model'),
G
Guanghua Yu 已提交
1070 1071
                input_spec=pruned_input_spec)
        logger.info("Export model and saved in {}".format(save_dir))
1072

G
Guanghua Yu 已提交
1073 1074 1075 1076 1077 1078 1079 1080 1081 1082 1083 1084
    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 已提交
1085
        kl_quant = True if hasattr(self.cfg.slim, 'ptq') else False
G
Guanghua Yu 已提交
1086
        _, pruned_input_spec = self._get_infer_cfg_and_input_spec(
S
shangliang Xu 已提交
1087
            save_dir, prune_input=False, kl_quant=kl_quant)
G
Guanghua Yu 已提交
1088 1089 1090 1091 1092 1093

        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 已提交
1094 1095 1096 1097 1098 1099 1100 1101 1102 1103 1104 1105 1106 1107 1108 1109 1110 1111 1112 1113 1114 1115 1116 1117 1118

    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))
1119 1120 1121 1122 1123 1124 1125 1126 1127 1128 1129 1130 1131 1132 1133 1134 1135 1136 1137 1138 1139 1140 1141

    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)
1142 1143 1144
            logger.info("Found {} inference images in total.".format(
                len(images)))
        return all_images