AIStoryBoard/python/PaddleOCR-main/benchmark/PaddleOCR_DBNet/trainer/trainer.py

# -*- coding: utf-8 -*-
# @Time    : 2019/8/23 21:58
# @Author  : zhoujun
import time

import paddle
from tqdm import tqdm

from base import BaseTrainer
from utils import runningScore, cal_text_score, Polynomial, profiler


class Trainer(BaseTrainer):
    def __init__(
        self,
        config,
        model,
        criterion,
        train_loader,
        validate_loader,
        metric_cls,
        post_process=None,
        profiler_options=None,
    ):
        super(Trainer, self).__init__(
            config,
            model,
            criterion,
            train_loader,
            validate_loader,
            metric_cls,
            post_process,
        )
        self.profiler_options = profiler_options
        self.enable_eval = config["trainer"].get("enable_eval", True)

    def _train_epoch(self, epoch):
        self.model.train()
        total_samples = 0
        train_reader_cost = 0.0
        train_batch_cost = 0.0
        reader_start = time.time()
        epoch_start = time.time()
        train_loss = 0.0
        running_metric_text = runningScore(2)

        for i, batch in enumerate(self.train_loader):
            profiler.add_profiler_step(self.profiler_options)
            if i >= self.train_loader_len:
                break
            self.global_step += 1
            lr = self.optimizer.get_lr()

            cur_batch_size = batch["img"].shape[0]

            train_reader_cost += time.time() - reader_start
            if self.amp:
                with paddle.amp.auto_cast(
                    enable="gpu" in paddle.device.get_device(),
                    custom_white_list=self.amp.get("custom_white_list", []),
                    custom_black_list=self.amp.get("custom_black_list", []),
                    level=self.amp.get("level", "O2"),
                ):
                    preds = self.model(batch["img"])
                loss_dict = self.criterion(preds.astype(paddle.float32), batch)
                scaled_loss = self.amp["scaler"].scale(loss_dict["loss"])
                scaled_loss.backward()
                self.amp["scaler"].minimize(self.optimizer, scaled_loss)
            else:
                preds = self.model(batch["img"])
                loss_dict = self.criterion(preds, batch)
                # backward
                loss_dict["loss"].backward()
                self.optimizer.step()
            self.lr_scheduler.step()
            self.optimizer.clear_grad()

            train_batch_time = time.time() - reader_start
            train_batch_cost += train_batch_time
            total_samples += cur_batch_size

            # acc iou
            score_shrink_map = cal_text_score(
                preds[:, 0, :, :],
                batch["shrink_map"],
                batch["shrink_mask"],
                running_metric_text,
                thred=self.config["post_processing"]["args"]["thresh"],
            )

            # loss 和 acc 记录到日志
            loss_str = "loss: {:.4f}, ".format(loss_dict["loss"].item())
            for idx, (key, value) in enumerate(loss_dict.items()):
                loss_dict[key] = value.item()
                if key == "loss":
                    continue
                loss_str += "{}: {:.4f}".format(key, loss_dict[key])
                if idx < len(loss_dict) - 1:
                    loss_str += ", "

            train_loss += loss_dict["loss"]
            acc = score_shrink_map["Mean Acc"]
            iou_shrink_map = score_shrink_map["Mean IoU"]

            if self.global_step % self.log_iter == 0:
                self.logger_info(
                    "[{}/{}], [{}/{}], global_step: {}, ips: {:.1f} samples/sec, avg_reader_cost: {:.5f} s, avg_batch_cost: {:.5f} s, avg_samples: {}, acc: {:.4f}, iou_shrink_map: {:.4f}, {}lr:{:.6}, time:{:.2f}".format(
                        epoch,
                        self.epochs,
                        i + 1,
                        self.train_loader_len,
                        self.global_step,
                        total_samples / train_batch_cost,
                        train_reader_cost / self.log_iter,
                        train_batch_cost / self.log_iter,
                        total_samples / self.log_iter,
                        acc,
                        iou_shrink_map,
                        loss_str,
                        lr,
                        train_batch_cost,
                    )
                )
                total_samples = 0
                train_reader_cost = 0.0
                train_batch_cost = 0.0

            if self.visualdl_enable and paddle.distributed.get_rank() == 0:
                # write tensorboard
                for key, value in loss_dict.items():
                    self.writer.add_scalar(
                        "TRAIN/LOSS/{}".format(key), value, self.global_step
                    )
                self.writer.add_scalar("TRAIN/ACC_IOU/acc", acc, self.global_step)
                self.writer.add_scalar(
                    "TRAIN/ACC_IOU/iou_shrink_map", iou_shrink_map, self.global_step
                )
                self.writer.add_scalar("TRAIN/lr", lr, self.global_step)
            reader_start = time.time()
        return {
            "train_loss": train_loss / self.train_loader_len,
            "lr": lr,
            "time": time.time() - epoch_start,
            "epoch": epoch,
        }

    def _eval(self, epoch):
        self.model.eval()
        raw_metrics = []
        total_frame = 0.0
        total_time = 0.0
        for i, batch in tqdm(
            enumerate(self.validate_loader),
            total=len(self.validate_loader),
            desc="test model",
        ):
            with paddle.no_grad():
                start = time.time()
                if self.amp:
                    with paddle.amp.auto_cast(
                        enable="gpu" in paddle.device.get_device(),
                        custom_white_list=self.amp.get("custom_white_list", []),
                        custom_black_list=self.amp.get("custom_black_list", []),
                        level=self.amp.get("level", "O2"),
                    ):
                        preds = self.model(batch["img"])
                    preds = preds.astype(paddle.float32)
                else:
                    preds = self.model(batch["img"])
                boxes, scores = self.post_process(
                    batch, preds, is_output_polygon=self.metric_cls.is_output_polygon
                )
                total_frame += batch["img"].shape[0]
                total_time += time.time() - start
                raw_metric = self.metric_cls.validate_measure(batch, (boxes, scores))
                raw_metrics.append(raw_metric)
        metrics = self.metric_cls.gather_measure(raw_metrics)
        self.logger_info("FPS:{}".format(total_frame / total_time))
        return metrics["recall"].avg, metrics["precision"].avg, metrics["fmeasure"].avg

    def _on_epoch_finish(self):
        self.logger_info(
            "[{}/{}], train_loss: {:.4f}, time: {:.4f}, lr: {}".format(
                self.epoch_result["epoch"],
                self.epochs,
                self.epoch_result["train_loss"],
                self.epoch_result["time"],
                self.epoch_result["lr"],
            )
        )
        net_save_path = "{}/model_latest.pth".format(self.checkpoint_dir)
        net_save_path_best = "{}/model_best.pth".format(self.checkpoint_dir)

        if paddle.distributed.get_rank() == 0:
            self._save_checkpoint(self.epoch_result["epoch"], net_save_path)
            save_best = False
            if (
                self.validate_loader is not None
                and self.metric_cls is not None
                and self.enable_eval
            ):  # 使用f1作为最优模型指标
                recall, precision, hmean = self._eval(self.epoch_result["epoch"])

                if self.visualdl_enable:
                    self.writer.add_scalar("EVAL/recall", recall, self.global_step)
                    self.writer.add_scalar(
                        "EVAL/precision", precision, self.global_step
                    )
                    self.writer.add_scalar("EVAL/hmean", hmean, self.global_step)
                self.logger_info(
                    "test: recall: {:.6f}, precision: {:.6f}, hmean: {:.6f}".format(
                        recall, precision, hmean
                    )
                )

                if hmean >= self.metrics["hmean"]:
                    save_best = True
                    self.metrics["train_loss"] = self.epoch_result["train_loss"]
                    self.metrics["hmean"] = hmean
                    self.metrics["precision"] = precision
                    self.metrics["recall"] = recall
                    self.metrics["best_model_epoch"] = self.epoch_result["epoch"]
            else:
                if self.epoch_result["train_loss"] <= self.metrics["train_loss"]:
                    save_best = True
                    self.metrics["train_loss"] = self.epoch_result["train_loss"]
                    self.metrics["best_model_epoch"] = self.epoch_result["epoch"]
            best_str = "current best, "
            for k, v in self.metrics.items():
                best_str += "{}: {:.6f}, ".format(k, v)
            self.logger_info(best_str)
            if save_best:
                import shutil

                shutil.copy(net_save_path, net_save_path_best)
                self.logger_info("Saving current best: {}".format(net_save_path_best))
            else:
                self.logger_info("Saving checkpoint: {}".format(net_save_path))

    def _on_train_finish(self):
        if self.enable_eval:
            for k, v in self.metrics.items():
                self.logger_info("{}:{}".format(k, v))
        self.logger_info("finish train")

    def _initialize_scheduler(self):
        if self.config["lr_scheduler"]["type"] == "Polynomial":
            self.config["lr_scheduler"]["args"]["epochs"] = self.config["trainer"][
                "epochs"
            ]
            self.config["lr_scheduler"]["args"]["step_each_epoch"] = len(
                self.train_loader
            )
            self.lr_scheduler = Polynomial(**self.config["lr_scheduler"]["args"])()
        else:
            self.lr_scheduler = self._initialize("lr_scheduler", paddle.optimizer.lr)