| |
| |
| |
|
|
| import os |
| import pathlib |
| import shutil |
| from pprint import pformat |
|
|
| import anyconfig |
| import paddle |
| import numpy as np |
| import random |
| from paddle.jit import to_static |
| from paddle.static import InputSpec |
|
|
| from utils import setup_logger |
|
|
|
|
| class BaseTrainer: |
| def __init__( |
| self, |
| config, |
| model, |
| criterion, |
| train_loader, |
| validate_loader, |
| metric_cls, |
| post_process=None, |
| ): |
| config["trainer"]["output_dir"] = os.path.join( |
| str(pathlib.Path(os.path.abspath(__name__)).parent), |
| config["trainer"]["output_dir"], |
| ) |
| config["name"] = config["name"] + "_" + model.name |
| self.save_dir = config["trainer"]["output_dir"] |
| self.checkpoint_dir = os.path.join(self.save_dir, "checkpoint") |
|
|
| os.makedirs(self.checkpoint_dir, exist_ok=True) |
|
|
| self.global_step = 0 |
| self.start_epoch = 0 |
| self.config = config |
| self.criterion = criterion |
| |
| self.visualdl_enable = self.config["trainer"].get("visual_dl", False) |
| self.epochs = self.config["trainer"]["epochs"] |
| self.log_iter = self.config["trainer"]["log_iter"] |
| if paddle.distributed.get_rank() == 0: |
| anyconfig.dump(config, os.path.join(self.save_dir, "config.yaml")) |
| self.logger = setup_logger(os.path.join(self.save_dir, "train.log")) |
| self.logger_info(pformat(self.config)) |
|
|
| self.model = self.apply_to_static(model) |
|
|
| |
| if ( |
| paddle.device.cuda.device_count() > 0 |
| and paddle.device.is_compiled_with_cuda() |
| ): |
| self.with_cuda = True |
| random.seed(self.config["trainer"]["seed"]) |
| np.random.seed(self.config["trainer"]["seed"]) |
| paddle.seed(self.config["trainer"]["seed"]) |
| else: |
| self.with_cuda = False |
| self.logger_info("train with and paddle {}".format(paddle.__version__)) |
| |
| self.metrics = { |
| "recall": 0, |
| "precision": 0, |
| "hmean": 0, |
| "train_loss": float("inf"), |
| "best_model_epoch": 0, |
| } |
|
|
| self.train_loader = train_loader |
| if validate_loader is not None: |
| assert post_process is not None and metric_cls is not None |
| self.validate_loader = validate_loader |
| self.post_process = post_process |
| self.metric_cls = metric_cls |
| self.train_loader_len = len(train_loader) |
|
|
| if self.validate_loader is not None: |
| self.logger_info( |
| "train dataset has {} samples,{} in dataloader, validate dataset has {} samples,{} in dataloader".format( |
| len(self.train_loader.dataset), |
| self.train_loader_len, |
| len(self.validate_loader.dataset), |
| len(self.validate_loader), |
| ) |
| ) |
| else: |
| self.logger_info( |
| "train dataset has {} samples,{} in dataloader".format( |
| len(self.train_loader.dataset), self.train_loader_len |
| ) |
| ) |
|
|
| self._initialize_scheduler() |
|
|
| self._initialize_optimizer() |
|
|
| |
| if self.config["trainer"]["resume_checkpoint"] != "": |
| self._load_checkpoint( |
| self.config["trainer"]["resume_checkpoint"], resume=True |
| ) |
| elif self.config["trainer"]["finetune_checkpoint"] != "": |
| self._load_checkpoint( |
| self.config["trainer"]["finetune_checkpoint"], resume=False |
| ) |
|
|
| if self.visualdl_enable and paddle.distributed.get_rank() == 0: |
| from visualdl import LogWriter |
|
|
| self.writer = LogWriter(self.save_dir) |
|
|
| |
| self.amp = self.config.get("amp", None) |
| if self.amp == "None": |
| self.amp = None |
| if self.amp: |
| self.amp["scaler"] = paddle.amp.GradScaler( |
| init_loss_scaling=self.amp.get("scale_loss", 1024), |
| use_dynamic_loss_scaling=self.amp.get("use_dynamic_loss_scaling", True), |
| ) |
| self.model, self.optimizer = paddle.amp.decorate( |
| models=self.model, |
| optimizers=self.optimizer, |
| level=self.amp.get("amp_level", "O2"), |
| ) |
|
|
| |
| if paddle.device.cuda.device_count() > 1: |
| self.model = paddle.DataParallel(self.model) |
| |
| self.UN_Normalize = False |
| for t in self.config["dataset"]["train"]["dataset"]["args"]["transforms"]: |
| if t["type"] == "Normalize": |
| self.normalize_mean = t["args"]["mean"] |
| self.normalize_std = t["args"]["std"] |
| self.UN_Normalize = True |
|
|
| def apply_to_static(self, model): |
| support_to_static = self.config["trainer"].get("to_static", False) |
| if support_to_static: |
| specs = None |
| print("static") |
| specs = [InputSpec([None, 3, -1, -1])] |
| model = to_static(model, input_spec=specs) |
| self.logger_info( |
| "Successfully to apply @to_static with specs: {}".format(specs) |
| ) |
| return model |
|
|
| def train(self): |
| """ |
| Full training logic |
| """ |
| for epoch in range(self.start_epoch + 1, self.epochs + 1): |
| self.epoch_result = self._train_epoch(epoch) |
| self._on_epoch_finish() |
| if paddle.distributed.get_rank() == 0 and self.visualdl_enable: |
| self.writer.close() |
| self._on_train_finish() |
|
|
| def _train_epoch(self, epoch): |
| """ |
| Training logic for an epoch |
| |
| :param epoch: Current epoch number |
| """ |
| raise NotImplementedError |
|
|
| def _eval(self, epoch): |
| """ |
| eval logic for an epoch |
| |
| :param epoch: Current epoch number |
| """ |
| raise NotImplementedError |
|
|
| def _on_epoch_finish(self): |
| raise NotImplementedError |
|
|
| def _on_train_finish(self): |
| raise NotImplementedError |
|
|
| def _save_checkpoint(self, epoch, file_name): |
| """ |
| Saving checkpoints |
| |
| :param epoch: current epoch number |
| :param log: logging information of the epoch |
| :param save_best: if True, rename the saved checkpoint to 'model_best.pth.tar' |
| """ |
| state_dict = self.model.state_dict() |
| state = { |
| "epoch": epoch, |
| "global_step": self.global_step, |
| "state_dict": state_dict, |
| "optimizer": self.optimizer.state_dict(), |
| "config": self.config, |
| "metrics": self.metrics, |
| } |
| filename = os.path.join(self.checkpoint_dir, file_name) |
| paddle.save(state, filename) |
|
|
| def _load_checkpoint(self, checkpoint_path, resume): |
| """ |
| Resume from saved checkpoints |
| :param checkpoint_path: Checkpoint path to be resumed |
| """ |
| self.logger_info("Loading checkpoint: {} ...".format(checkpoint_path)) |
| checkpoint = paddle.load(checkpoint_path) |
| self.model.set_state_dict(checkpoint["state_dict"]) |
| if resume: |
| self.global_step = checkpoint["global_step"] |
| self.start_epoch = checkpoint["epoch"] |
| self.config["lr_scheduler"]["args"]["last_epoch"] = self.start_epoch |
| |
| self.optimizer.set_state_dict(checkpoint["optimizer"]) |
| if "metrics" in checkpoint: |
| self.metrics = checkpoint["metrics"] |
| self.logger_info( |
| "resume from checkpoint {} (epoch {})".format( |
| checkpoint_path, self.start_epoch |
| ) |
| ) |
| else: |
| self.logger_info("finetune from checkpoint {}".format(checkpoint_path)) |
|
|
| def _initialize(self, name, module, *args, **kwargs): |
| module_name = self.config[name]["type"] |
| module_args = self.config[name].get("args", {}) |
| assert all( |
| [k not in module_args for k in kwargs] |
| ), "Overwriting kwargs given in config file is not allowed" |
| module_args.update(kwargs) |
| return getattr(module, module_name)(*args, **module_args) |
|
|
| def _initialize_scheduler(self): |
| self.lr_scheduler = self._initialize("lr_scheduler", paddle.optimizer.lr) |
|
|
| def _initialize_optimizer(self): |
| self.optimizer = self._initialize( |
| "optimizer", |
| paddle.optimizer, |
| parameters=self.model.parameters(), |
| learning_rate=self.lr_scheduler, |
| ) |
|
|
| def inverse_normalize(self, batch_img): |
| if self.UN_Normalize: |
| batch_img[:, 0, :, :] = ( |
| batch_img[:, 0, :, :] * self.normalize_std[0] + self.normalize_mean[0] |
| ) |
| batch_img[:, 1, :, :] = ( |
| batch_img[:, 1, :, :] * self.normalize_std[1] + self.normalize_mean[1] |
| ) |
| batch_img[:, 2, :, :] = ( |
| batch_img[:, 2, :, :] * self.normalize_std[2] + self.normalize_mean[2] |
| ) |
|
|
| def logger_info(self, s): |
| if paddle.distributed.get_rank() == 0: |
| self.logger.info(s) |
|
|
