# 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. import os import time from collections import deque, defaultdict import pickle import shutil import numpy as np import paddle import paddle.nn.functional as F from paddleseg.utils import TimeAverager, calculate_eta, resume, logger from .val import evaluate def visual_in_traning(log_writer, vis_dict, step): """ Visual in vdl Args: log_writer (LogWriter): The log writer of vdl. vis_dict (dict): Dict of tensor. The shape of thesor is (C, H, W) """ for key, value in vis_dict.items(): value_shape = value.shape if value_shape[0] not in [1, 3]: value = value[0] value = value.unsqueeze(0) value = paddle.transpose(value, (1, 2, 0)) min_v = paddle.min(value) max_v = paddle.max(value) if (min_v > 0) and (max_v < 1): value = value * 255 elif (min_v < 0 and min_v >= -1) and (max_v <= 1): value = (1 + value) / 2 * 255 else: value = (value - min_v) / (max_v - min_v) * 255 value = value.astype('uint8') value = value.numpy() log_writer.add_image(tag=key, img=value, step=step) def save_best(best_model_dir, metrics_data, iter): with open(os.path.join(best_model_dir, 'best_metrics.txt'), 'w') as f: for key, value in metrics_data.items(): line = key + ' ' + str(value) + '\n' f.write(line) f.write('iter' + ' ' + str(iter) + '\n') def get_best(best_file, metrics, resume_model=None): '''Get best metrics and iter from file''' best_metrics_data = {} if os.path.exists(best_file) and (resume_model is not None): values = [] with open(best_file, 'r') as f: lines = f.readlines() for line in lines: line = line.strip() key, value = line.split(' ') best_metrics_data[key] = eval(value) if key == 'iter': best_iter = eval(value) else: for key in metrics: best_metrics_data[key] = np.inf best_iter = -1 return best_metrics_data, best_iter def train(model, train_dataset, val_dataset=None, optimizer=None, save_dir='output', iters=10000, batch_size=2, resume_model=None, save_interval=1000, log_iters=10, log_image_iters=1000, num_workers=0, use_vdl=False, losses=None, keep_checkpoint_max=5, eval_begin_iters=None, metrics='sad'): """ Launch training. Args: model（nn.Layer): A matting model. train_dataset (paddle.io.Dataset): Used to read and process training datasets. val_dataset (paddle.io.Dataset, optional): Used to read and process validation datasets. optimizer (paddle.optimizer.Optimizer): The optimizer. save_dir (str, optional): The directory for saving the model snapshot. Default: 'output'. iters (int, optional): How may iters to train the model. Defualt: 10000. batch_size (int, optional): Mini batch size of one gpu or cpu. Default: 2. resume_model (str, optional): The path of resume model. save_interval (int, optional): How many iters to save a model snapshot once during training. Default: 1000. log_iters (int, optional): Display logging information at every log_iters. Default: 10. log_image_iters (int, optional): Log image to vdl. Default: 1000. num_workers (int, optional): Num workers for data loader. Default: 0. use_vdl (bool, optional): Whether to record the data to VisualDL during training. Default: False. losses (dict, optional): A dict of loss, refer to the loss function of the model for details. Default: None. keep_checkpoint_max (int, optional): Maximum number of checkpoints to save. Default: 5. eval_begin_iters (int): The iters begin evaluation. It will evaluate at iters/2 if it is None. Defalust: None. metrics(str|list, optional): The metrics to evaluate, it may be the combination of ("sad", "mse", "grad", "conn"). """ model.train() nranks = paddle.distributed.ParallelEnv().nranks local_rank = paddle.distributed.ParallelEnv().local_rank start_iter = 0 if resume_model is not None: start_iter = resume(model, optimizer, resume_model) if not os.path.isdir(save_dir): if os.path.exists(save_dir): os.remove(save_dir) os.makedirs(save_dir) if nranks > 1: # Initialize parallel environment if not done. if not paddle.distributed.parallel.parallel_helper._is_parallel_ctx_initialized( ): paddle.distributed.init_parallel_env() ddp_model = paddle.DataParallel(model) else: ddp_model = paddle.DataParallel(model) batch_sampler = paddle.io.DistributedBatchSampler( train_dataset, batch_size=batch_size, shuffle=True, drop_last=True) loader = paddle.io.DataLoader( train_dataset, batch_sampler=batch_sampler, num_workers=num_workers, return_list=True, ) if use_vdl: from visualdl import LogWriter log_writer = LogWriter(save_dir) if isinstance(metrics, str): metrics = [metrics] elif not isinstance(metrics, list): metrics = ['sad'] best_metrics_data, best_iter = get_best( os.path.join(save_dir, 'best_model', 'best_metrics.txt'), metrics, resume_model=resume_model) avg_loss = defaultdict(float) iters_per_epoch = len(batch_sampler) reader_cost_averager = TimeAverager() batch_cost_averager = TimeAverager() save_models = deque() batch_start = time.time() iter = start_iter while iter < iters: for data in loader: iter += 1 if iter > iters: break reader_cost_averager.record(time.time() - batch_start) logit_dict, loss_dict = ddp_model(data) if nranks > 1 else model( data) loss_dict['all'].backward() optimizer.step() lr = optimizer.get_lr() if isinstance(optimizer._learning_rate, paddle.optimizer.lr.LRScheduler): optimizer._learning_rate.step() model.clear_gradients() for key, value in loss_dict.items(): avg_loss[key] += value.numpy()[0] batch_cost_averager.record( time.time() - batch_start, num_samples=batch_size) if (iter) % log_iters == 0 and local_rank == 0: for key, value in avg_loss.items(): avg_loss[key] = value / log_iters remain_iters = iters - iter avg_train_batch_cost = batch_cost_averager.get_average() avg_train_reader_cost = reader_cost_averager.get_average() eta = calculate_eta(remain_iters, avg_train_batch_cost) # loss info loss_str = ' ' * 26 + '\t[LOSSES]' loss_str = loss_str for key, value in avg_loss.items(): if key != 'all': loss_str = loss_str + ' ' + key + '={:.4f}'.format( value) logger.info( "[TRAIN] epoch={}, iter={}/{}, loss={:.4f}, lr={:.6f}, batch_cost={:.4f}, reader_cost={:.5f}, ips={:.4f} samples/sec | ETA {}\n{}\n" .format((iter - 1) // iters_per_epoch + 1, iter, iters, avg_loss['all'], lr, avg_train_batch_cost, avg_train_reader_cost, batch_cost_averager.get_ips_average( ), eta, loss_str)) if use_vdl: for key, value in avg_loss.items(): log_tag = 'Train/' + key log_writer.add_scalar(log_tag, value, iter) log_writer.add_scalar('Train/lr', lr, iter) log_writer.add_scalar('Train/batch_cost', avg_train_batch_cost, iter) log_writer.add_scalar('Train/reader_cost', avg_train_reader_cost, iter) if iter % log_image_iters == 0: vis_dict = {} # ground truth vis_dict['ground truth/img'] = data['img'][0] for key in data['gt_fields']: key = key[0] vis_dict['/'.join(['ground truth', key])] = data[ key][0] # predict for key, value in logit_dict.items(): vis_dict['/'.join(['predict', key])] = logit_dict[ key][0] visual_in_traning( log_writer=log_writer, vis_dict=vis_dict, step=iter) for key in avg_loss.keys(): avg_loss[key] = 0. reader_cost_averager.reset() batch_cost_averager.reset() # save model if (iter % save_interval == 0 or iter == iters) and local_rank == 0: current_save_dir = os.path.join(save_dir, "iter_{}".format(iter)) if not os.path.isdir(current_save_dir): os.makedirs(current_save_dir) paddle.save(model.state_dict(), os.path.join(current_save_dir, 'model.pdparams')) paddle.save(optimizer.state_dict(), os.path.join(current_save_dir, 'model.pdopt')) save_models.append(current_save_dir) if len(save_models) > keep_checkpoint_max > 0: model_to_remove = save_models.popleft() shutil.rmtree(model_to_remove) # eval model if eval_begin_iters is None: eval_begin_iters = iters // 2 if (iter % save_interval == 0 or iter == iters) and ( val_dataset is not None ) and local_rank == 0 and iter >= eval_begin_iters: num_workers = 1 if num_workers > 0 else 0 metrics_data = evaluate( model, val_dataset, num_workers=1, print_detail=True, save_results=False, metrics=metrics) model.train() # save best model and add evaluation results to vdl if (iter % save_interval == 0 or iter == iters) and local_rank == 0: if val_dataset is not None and iter >= eval_begin_iters: if metrics_data[metrics[0]] < best_metrics_data[metrics[0]]: best_iter = iter best_metrics_data = metrics_data.copy() best_model_dir = os.path.join(save_dir, "best_model") paddle.save( model.state_dict(), os.path.join(best_model_dir, 'model.pdparams')) save_best(best_model_dir, best_metrics_data, iter) show_list = [] for key, value in best_metrics_data.items(): show_list.append((key, value)) log_str = '[EVAL] The model with the best validation {} ({:.4f}) was saved at iter {}.'.format( show_list[0][0], show_list[0][1], best_iter) if len(show_list) > 1: log_str += " While" for i in range(1, len(show_list)): log_str = log_str + ' {}: {:.4f},'.format( show_list[i][0], show_list[i][1]) log_str = log_str[:-1] logger.info(log_str) if use_vdl: for key, value in metrics_data.items(): log_writer.add_scalar('Evaluate/' + key, value, iter) batch_start = time.time() # Sleep for half a second to let dataloader release resources. time.sleep(0.5) if use_vdl: log_writer.close()