init
This commit is contained in:
@@ -0,0 +1,2 @@
|
||||
from .base_trainer import BaseTrainer
|
||||
from .base_dataset import BaseDataSet
|
||||
@@ -0,0 +1,87 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# @Time : 2019/12/4 13:12
|
||||
# @Author : zhoujun
|
||||
import copy
|
||||
from paddle.io import Dataset
|
||||
from data_loader.modules import *
|
||||
|
||||
|
||||
class BaseDataSet(Dataset):
|
||||
def __init__(self,
|
||||
data_path: str,
|
||||
img_mode,
|
||||
pre_processes,
|
||||
filter_keys,
|
||||
ignore_tags,
|
||||
transform=None,
|
||||
target_transform=None):
|
||||
assert img_mode in ['RGB', 'BRG', 'GRAY']
|
||||
self.ignore_tags = ignore_tags
|
||||
self.data_list = self.load_data(data_path)
|
||||
item_keys = [
|
||||
'img_path', 'img_name', 'text_polys', 'texts', 'ignore_tags'
|
||||
]
|
||||
for item in item_keys:
|
||||
assert item in self.data_list[
|
||||
0], 'data_list from load_data must contains {}'.format(
|
||||
item_keys)
|
||||
self.img_mode = img_mode
|
||||
self.filter_keys = filter_keys
|
||||
self.transform = transform
|
||||
self.target_transform = target_transform
|
||||
self._init_pre_processes(pre_processes)
|
||||
|
||||
def _init_pre_processes(self, pre_processes):
|
||||
self.aug = []
|
||||
if pre_processes is not None:
|
||||
for aug in pre_processes:
|
||||
if 'args' not in aug:
|
||||
args = {}
|
||||
else:
|
||||
args = aug['args']
|
||||
if isinstance(args, dict):
|
||||
cls = eval(aug['type'])(**args)
|
||||
else:
|
||||
cls = eval(aug['type'])(args)
|
||||
self.aug.append(cls)
|
||||
|
||||
def load_data(self, data_path: str) -> list:
|
||||
"""
|
||||
把数据加载为一个list：
|
||||
:params data_path: 存储数据的文件夹或者文件
|
||||
return a dict ,包含了，'img_path','img_name','text_polys','texts','ignore_tags'
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def apply_pre_processes(self, data):
|
||||
for aug in self.aug:
|
||||
data = aug(data)
|
||||
return data
|
||||
|
||||
def __getitem__(self, index):
|
||||
try:
|
||||
data = copy.deepcopy(self.data_list[index])
|
||||
im = cv2.imread(data['img_path'], 1
|
||||
if self.img_mode != 'GRAY' else 0)
|
||||
if self.img_mode == 'RGB':
|
||||
im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
|
||||
data['img'] = im
|
||||
data['shape'] = [im.shape[0], im.shape[1]]
|
||||
data = self.apply_pre_processes(data)
|
||||
|
||||
if self.transform:
|
||||
data['img'] = self.transform(data['img'])
|
||||
data['text_polys'] = data['text_polys'].tolist()
|
||||
if len(self.filter_keys):
|
||||
data_dict = {}
|
||||
for k, v in data.items():
|
||||
if k not in self.filter_keys:
|
||||
data_dict[k] = v
|
||||
return data_dict
|
||||
else:
|
||||
return data
|
||||
except:
|
||||
return self.__getitem__(np.random.randint(self.__len__()))
|
||||
|
||||
def __len__(self):
|
||||
return len(self.data_list)
|
||||
@@ -0,0 +1,250 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# @Time : 2019/8/23 21:50
|
||||
# @Author : zhoujun
|
||||
|
||||
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
|
||||
# logger and tensorboard
|
||||
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)
|
||||
|
||||
# device
|
||||
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__))
|
||||
# metrics
|
||||
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()
|
||||
|
||||
# resume or finetune
|
||||
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)
|
||||
# make inverse Normalize
|
||||
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.scheduler.load_state_dict(checkpoint['scheduler'])
|
||||
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)
|
||||
Reference in New Issue
Block a user