train.py 8.7 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22
# Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

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

import os
import time
import multiprocessing
import numpy as np
23 24
import datetime
from collections import deque
25

26 27 28 29 30
def set_paddle_flags(**kwargs):
    for key, value in kwargs.items():
        if os.environ.get(key, None) is None:
            os.environ[key] = str(value)

31 32
# NOTE(paddle-dev): All of these flags should be set before 
# `import paddle`. Otherwise, it would not take any effect.
33 34 35 36
set_paddle_flags(
    FLAGS_eager_delete_tensor_gb=0,  # enable GC to save memory
)

37 38 39 40 41 42 43
from paddle import fluid

from ppdet.core.workspace import load_config, merge_config, create
from ppdet.data.data_feed import create_reader

from ppdet.utils.eval_utils import parse_fetches, eval_run, eval_results
from ppdet.utils.stats import TrainingStats
Y
Yang Zhang 已提交
44
from ppdet.utils.cli import ArgsParser
45
from ppdet.utils.check import check_gpu
46
import ppdet.utils.checkpoint as checkpoint
47
from ppdet.modeling.model_input import create_feed
48 49 50 51 52 53 54 55

import logging
FORMAT = '%(asctime)s-%(levelname)s: %(message)s'
logging.basicConfig(level=logging.INFO, format=FORMAT)
logger = logging.getLogger(__name__)


def main():
Y
Yang Zhang 已提交
56
    cfg = load_config(FLAGS.config)
57
    if 'architecture' in cfg:
Y
Yang Zhang 已提交
58
        main_arch = cfg.architecture
59 60 61
    else:
        raise ValueError("'architecture' not specified in config file.")

Y
Yang Zhang 已提交
62
    merge_config(FLAGS.opt)
63 64
    if 'log_iter' not in cfg:
        cfg.log_iter = 20
65

66 67 68
    # check if set use_gpu=True in paddlepaddle cpu version
    check_gpu(cfg.use_gpu)

Y
Yang Zhang 已提交
69
    if cfg.use_gpu:
70 71
        devices_num = fluid.core.get_cuda_device_count()
    else:
72 73
        devices_num = int(
            os.environ.get('CPU_NUM', multiprocessing.cpu_count()))
74 75

    if 'train_feed' not in cfg:
76
        train_feed = create(main_arch + 'TrainFeed')
77
    else:
Y
Yang Zhang 已提交
78
        train_feed = create(cfg.train_feed)
79

Y
Yang Zhang 已提交
80
    if FLAGS.eval:
81
        if 'eval_feed' not in cfg:
82
            eval_feed = create(main_arch + 'EvalFeed')
83
        else:
Y
Yang Zhang 已提交
84
            eval_feed = create(cfg.eval_feed)
85

Y
Yang Zhang 已提交
86
    place = fluid.CUDAPlace(0) if cfg.use_gpu else fluid.CPUPlace()
87 88 89 90 91
    exe = fluid.Executor(place)

    lr_builder = create('LearningRate')
    optim_builder = create('OptimizerBuilder')

92
    # build program
93 94 95 96
    startup_prog = fluid.Program()
    train_prog = fluid.Program()
    with fluid.program_guard(train_prog, startup_prog):
        with fluid.unique_name.guard():
97
            model = create(main_arch)
98
            train_pyreader, feed_vars = create_feed(train_feed)
99 100 101 102 103 104
            train_fetches = model.train(feed_vars)
            loss = train_fetches['loss']
            lr = lr_builder()
            optimizer = optim_builder(lr)
            optimizer.minimize(loss)

W
wangguanzhong 已提交
105 106
    train_reader = create_reader(train_feed, cfg.max_iters * devices_num,
                                 FLAGS.dataset_dir)
107 108 109 110 111 112
    train_pyreader.decorate_sample_list_generator(train_reader, place)

    # parse train fetches
    train_keys, train_values, _ = parse_fetches(train_fetches)
    train_values.append(lr)

Y
Yang Zhang 已提交
113
    if FLAGS.eval:
114 115 116
        eval_prog = fluid.Program()
        with fluid.program_guard(eval_prog, startup_prog):
            with fluid.unique_name.guard():
117
                model = create(main_arch)
118
                eval_pyreader, feed_vars = create_feed(eval_feed)
119
                fetches = model.eval(feed_vars)
120 121
        eval_prog = eval_prog.clone(True)

W
wangguanzhong 已提交
122
        eval_reader = create_reader(eval_feed, args_path=FLAGS.dataset_dir)
123 124
        eval_pyreader.decorate_sample_list_generator(eval_reader, place)

125
        # parse eval fetches
126 127 128 129 130
        extra_keys = []
        if cfg.metric == 'COCO':
            extra_keys = ['im_info', 'im_id', 'im_shape']
        if cfg.metric == 'VOC':
            extra_keys = ['gt_box', 'gt_label', 'is_difficult']
131 132 133
        eval_keys, eval_values, eval_cls = parse_fetches(fetches, eval_prog,
                                                         extra_keys)

134
    # compile program for multi-devices
135 136
    build_strategy = fluid.BuildStrategy()
    build_strategy.memory_optimize = False
K
Kaipeng Deng 已提交
137
    build_strategy.enable_inplace = False
138
    sync_bn = getattr(model.backbone, 'norm_type', None) == 'sync_bn'
K
Kaipeng Deng 已提交
139 140
    # only enable sync_bn in multi GPU devices
    build_strategy.sync_batch_norm = sync_bn and devices_num > 1 \
141
         and cfg.use_gpu
142 143 144
    train_compile_program = fluid.compiler.CompiledProgram(
        train_prog).with_data_parallel(
            loss_name=loss.name, build_strategy=build_strategy)
Y
Yang Zhang 已提交
145
    if FLAGS.eval:
146 147 148 149
        eval_compile_program = fluid.compiler.CompiledProgram(eval_prog)

    exe.run(startup_prog)

150
    fuse_bn = getattr(model.backbone, 'norm_type', None) == 'affine_channel'
Q
qingqing01 已提交
151
    start_iter = 0
Y
Yang Zhang 已提交
152 153
    if FLAGS.resume_checkpoint:
        checkpoint.load_checkpoint(exe, train_prog, FLAGS.resume_checkpoint)
Q
qingqing01 已提交
154
        start_iter = checkpoint.global_step()
155
    elif cfg.pretrain_weights and fuse_bn:
Y
Yang Zhang 已提交
156 157 158 159
        checkpoint.load_and_fusebn(exe, train_prog, cfg.pretrain_weights)
    elif cfg.pretrain_weights:
        checkpoint.load_pretrain(exe, train_prog, cfg.pretrain_weights)

160 161 162 163 164 165
    # whether output bbox is normalized in model output layer
    is_bbox_normalized = False
    if hasattr(model, 'is_bbox_normalized') and \
            callable(model.is_bbox_normalized):
        is_bbox_normalized = model.is_bbox_normalized()

Y
Yang Zhang 已提交
166
    train_stats = TrainingStats(cfg.log_smooth_window, train_keys)
167 168 169 170
    train_pyreader.start()
    start_time = time.time()
    end_time = time.time()

Y
Yang Zhang 已提交
171 172
    cfg_name = os.path.basename(FLAGS.config).split('.')[0]
    save_dir = os.path.join(cfg.save_dir, cfg_name)
173
    time_stat = deque(maxlen=cfg.log_iter)
174
    best_box_ap_list = [0.0, 0]  #[map, iter]
Q
qingqing01 已提交
175
    for it in range(start_iter, cfg.max_iters):
176 177
        start_time = end_time
        end_time = time.time()
178 179 180 181
        time_stat.append(end_time - start_time)
        time_cost = np.mean(time_stat)
        eta_sec = (cfg.max_iters - it) * time_cost
        eta = str(datetime.timedelta(seconds=int(eta_sec)))
182 183 184 185
        outs = exe.run(train_compile_program, fetch_list=train_values)
        stats = {k: np.array(v).mean() for k, v in zip(train_keys, outs[:-1])}
        train_stats.update(stats)
        logs = train_stats.log()
186 187 188 189
        if it % cfg.log_iter == 0:
            strs = 'iter: {}, lr: {:.6f}, {}, time: {:.3f}, eta: {}'.format(
                it, np.mean(outs[-1]), logs, time_cost, eta)
            logger.info(strs)
190

191 192 193
        if it > 0 and it % cfg.snapshot_iter == 0 or it == cfg.max_iters - 1:
            save_name = str(it) if it != cfg.max_iters - 1 else "model_final"
            checkpoint.save(exe, train_prog, os.path.join(save_dir, save_name))
194

Y
Yang Zhang 已提交
195
            if FLAGS.eval:
196
                # evaluation
197 198
                results = eval_run(exe, eval_compile_program, eval_pyreader,
                                   eval_keys, eval_values, eval_cls)
Y
Yang Zhang 已提交
199 200 201
                resolution = None
                if 'mask' in results[0]:
                    resolution = model.mask_head.resolution
202
                box_ap_stats = eval_results(results, eval_feed, cfg.metric, cfg.num_classes,
203
                             resolution, is_bbox_normalized, FLAGS.output_eval)
204 205 206 207 208 209
                if box_ap_stats[0] > best_box_ap_list[0]:
                    best_box_ap_list[0] = box_ap_stats[0]
                    best_box_ap_list[1] = it
                    checkpoint.save(exe, train_prog, os.path.join(save_dir,"best_model"))
                logger.info("Best test box ap: {}, in iter: {}".format(
                    best_box_ap_list[0],best_box_ap_list[1]))
210 211 212 213 214

    train_pyreader.reset()


if __name__ == '__main__':
Y
Yang Zhang 已提交
215 216 217 218 219 220 221 222 223 224 225 226 227
    parser = ArgsParser()
    parser.add_argument(
        "-r",
        "--resume_checkpoint",
        default=None,
        type=str,
        help="Checkpoint path for resuming training.")
    parser.add_argument(
        "--eval",
        action='store_true',
        default=False,
        help="Whether to perform evaluation in train")
    parser.add_argument(
228
        "--output_eval",
Y
Yang Zhang 已提交
229 230
        default=None,
        type=str,
231
        help="Evaluation directory, default is current directory.")
W
wangguanzhong 已提交
232 233 234 235 236 237
    parser.add_argument(
        "-d",
        "--dataset_dir",
        default=None,
        type=str,
        help="Dataset path, same as DataFeed.dataset.dataset_dir")
Y
Yang Zhang 已提交
238
    FLAGS = parser.parse_args()
239
    main()