train.py 3.1 KB
Newer Older
D
dongshuilong 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18
# Copyright (c) 2021 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, division, print_function

import time
import paddle
from ppcls.engine.train.utils import update_loss, update_metric, log_info
D
dongshuilong 已提交
19
from ppcls.utils import profiler
D
dongshuilong 已提交
20 21


W
weishengyu 已提交
22
def train_epoch(engine, epoch_id, print_batch_step):
D
dongshuilong 已提交
23
    tic = time.time()
Z
zhangbo9674 已提交
24
    v_current = [int(i) for i in paddle.__version__.split(".")]
W
weishengyu 已提交
25 26
    for iter_id, batch in enumerate(engine.train_dataloader):
        if iter_id >= engine.max_iter:
D
dongshuilong 已提交
27
            break
D
dongshuilong 已提交
28
        profiler.add_profiler_step(engine.config["profiler_options"])
D
dongshuilong 已提交
29
        if iter_id == 5:
W
weishengyu 已提交
30 31 32 33
            for key in engine.time_info:
                engine.time_info[key].reset()
        engine.time_info["reader_cost"].update(time.time() - tic)
        if engine.use_dali:
D
dongshuilong 已提交
34 35 36 37 38
            batch = [
                paddle.to_tensor(batch[0]['data']),
                paddle.to_tensor(batch[0]['label'])
            ]
        batch_size = batch[0].shape[0]
39
        if not engine.config["Global"].get("use_multilabel", False):
G
gaotingquan 已提交
40
            batch[1] = batch[1].reshape([batch_size, -1])
W
weishengyu 已提交
41
        engine.global_step += 1
42

D
dongshuilong 已提交
43
        # image input
W
weishengyu 已提交
44
        if engine.amp:
45 46 47 48 49 50
            amp_level = engine.config['AMP'].get("level", "O1").upper()
            with paddle.amp.auto_cast(
                    custom_black_list={
                        "flatten_contiguous_range", "greater_than"
                    },
                    level=amp_level):
W
weishengyu 已提交
51
                out = forward(engine, batch)
52
                loss_dict = engine.train_loss_func(out, batch[1])
D
dongshuilong 已提交
53
        else:
W
weishengyu 已提交
54
            out = forward(engine, batch)
55
            loss_dict = engine.train_loss_func(out, batch[1])
D
dongshuilong 已提交
56 57

        # step opt and lr
W
weishengyu 已提交
58 59
        if engine.amp:
            scaled = engine.scaler.scale(loss_dict["loss"])
D
dongshuilong 已提交
60
            scaled.backward()
W
weishengyu 已提交
61
            engine.scaler.minimize(engine.optimizer, scaled)
D
dongshuilong 已提交
62 63
        else:
            loss_dict["loss"].backward()
W
weishengyu 已提交
64
            engine.optimizer.step()
Z
zhangbo9674 已提交
65
        engine.optimizer.clear_grad()
W
weishengyu 已提交
66
        engine.lr_sch.step()
D
dongshuilong 已提交
67 68 69

        # below code just for logging
        # update metric_for_logger
W
weishengyu 已提交
70
        update_metric(engine, out, batch, batch_size)
D
dongshuilong 已提交
71
        # update_loss_for_logger
W
weishengyu 已提交
72 73
        update_loss(engine, loss_dict, batch_size)
        engine.time_info["batch_cost"].update(time.time() - tic)
D
dongshuilong 已提交
74
        if iter_id % print_batch_step == 0:
W
weishengyu 已提交
75
            log_info(engine, batch_size, epoch_id, iter_id)
D
dongshuilong 已提交
76
        tic = time.time()
D
dongshuilong 已提交
77

D
dongshuilong 已提交
78

C
cuicheng01 已提交
79 80 81
def forward(engine, batch):
    if not engine.is_rec:
        return engine.model(batch[0])
D
dongshuilong 已提交
82
    else:
C
cuicheng01 已提交
83
        return engine.model(batch[0], batch[1])