train_dygraph.py

# Copyright (c) 2019 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

os.environ['FLAGS_eager_delete_tensor_gb'] = "0.0"
os.environ['FLAGS_fraction_of_gpu_memory_to_use'] = "0.99"

import paddle.fluid as fluid
import numpy as np
import paddle
import logging
import shutil
from datetime import datetime
from paddle.utils import Ploter

from danet import DANet
from options import Options
from utils.cityscapes_data import cityscapes_train
from utils.cityscapes_data import cityscapes_val
from utils.lr_scheduler import Lr


def get_model(args):
    model = DANet('DANet',
                  backbone=args.backbone,
                  num_classes=args.num_classes,
                  batch_size=args.batch_size,
                  dilated=args.dilated,
                  multi_grid=args.multi_grid,
                  multi_dilation=args.multi_dilation)
    return model


def mean_iou(pred, label, num_classes=19):
    label = fluid.layers.elementwise_min(fluid.layers.cast(label, np.int32),
                                         fluid.layers.assign(np.array([num_classes], dtype=np.int32)))
    label_ig = (label == num_classes).astype('int32')
    label_ng = (label != num_classes).astype('int32')
    pred = fluid.layers.cast(fluid.layers.argmax(pred, axis=1), 'int32')
    pred = pred * label_ng + label_ig * num_classes
    miou, wrong, correct = fluid.layers.mean_iou(pred, label, num_classes + 1)
    label.stop_gradient = True
    return miou, wrong, correct


def loss_fn(pred, pred2, pred3, label, num_classes=19):
    pred = fluid.layers.transpose(pred, perm=[0, 2, 3, 1])
    pred = fluid.layers.reshape(pred, [-1, num_classes])

    pred2 = fluid.layers.transpose(pred2, perm=[0, 2, 3, 1])
    pred2 = fluid.layers.reshape(pred2, [-1, num_classes])

    pred3 = fluid.layers.transpose(pred3, perm=[0, 2, 3, 1])
    pred3 = fluid.layers.reshape(pred3, [-1, num_classes])

    label = fluid.layers.reshape(label, [-1, 1])

    pred = fluid.layers.softmax(pred, use_cudnn=False)
    loss1 = fluid.layers.cross_entropy(pred, label, ignore_index=255)

    pred2 = fluid.layers.softmax(pred2, use_cudnn=False)
    loss2 = fluid.layers.cross_entropy(pred2, label, ignore_index=255)

    pred3 = fluid.layers.softmax(pred3, use_cudnn=False)
    loss3 = fluid.layers.cross_entropy(pred3, label, ignore_index=255)

    label.stop_gradient = True
    return loss1 + loss2 + loss3


def optimizer_setting(args):
    if args.weight_decay is not None:
        regular = fluid.regularizer.L2Decay(regularization_coeff=args.weight_decay)
    else:
        regular = None
    if args.lr_scheduler == 'poly':
        lr_scheduler = Lr(lr_policy='poly',
                          base_lr=args.lr,
                          epoch_nums=args.epoch_num,
                          step_per_epoch=args.step_per_epoch,
                          power=args.lr_pow,
                          warm_up=args.warm_up,
                          warmup_epoch=args.warmup_epoch)
        decayed_lr = lr_scheduler.get_lr()
    elif args.lr_scheduler == 'cosine':
        lr_scheduler = Lr(lr_policy='cosine',
                          base_lr=args.lr,
                          epoch_nums=args.epoch_num,
                          step_per_epoch=args.step_per_epoch,
                          warm_up=args.warm_up,
                          warmup_epoch=args.warmup_epoch)
        decayed_lr = lr_scheduler.get_lr()
    elif args.lr_scheduler == 'piecewise':
        lr_scheduler = Lr(lr_policy='piecewise',
                          base_lr=args.lr,
                          epoch_nums=args.epoch_num,
                          step_per_epoch=args.step_per_epoch,
                          warm_up=args.warm_up,
                          warmup_epoch=args.warmup_epoch,
                          decay_epoch=[50, 100, 150],
                          gamma=0.1)
        decayed_lr = lr_scheduler.get_lr()
    else:
        decayed_lr = args.lr
    return fluid.optimizer.MomentumOptimizer(learning_rate=decayed_lr,
                                             momentum=args.momentum,
                                             regularization=regular)


def main(args):

    batch_size = args.batch_size
    num_epochs = args.epoch_num
    num_classes = args.num_classes
    data_root = args.data_folder
    num = fluid.core.get_cuda_device_count()
    print('GPU设备数量： {}'.format(num))

    # program
    start_prog = fluid.default_startup_program()
    train_prog = fluid.default_main_program()

    start_prog.random_seed = args.seed
    train_prog.random_seed = args.seed

    logging.basicConfig(level=logging.INFO,
                        filename='DANet_{}_train.log'.format(args.backbone),
                        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    logging.info('DANet')
    logging.info(args)

    place = fluid.CUDAPlace(0) if args.cuda else fluid.CPUPlace()
    train_loss_title = 'Train_loss'
    test_loss_title = 'Test_loss'

    train_iou_title = 'Train_mIOU'
    test_iou_title = 'Test_mIOU'

    plot_loss = Ploter(train_loss_title, test_loss_title)
    plot_iou = Ploter(train_iou_title, test_iou_title)

    with fluid.dygraph.guard(place):

        model = get_model(args)
        x = np.random.randn(batch_size, 3, 224, 224).astype('float32')
        x = fluid.dygraph.to_variable(x)
        model(x)

        # 加载预训练模型
        if args.load_pretrained_model:
            save_dir = 'checkpoint/DANet101_pretrained_model_paddle1.6'
            if os.path.exists(save_dir + '.pdparams'):
                param, _ = fluid.load_dygraph(save_dir)
                model.set_dict(param)
                assert len(param) == len(model.state_dict()), "参数量不一致，加载参数失败，" \
                                                              "请核对模型是否初始化/模型是否一致"
                print('load pretrained model!')

        # 加载最优模型
        if args.load_better_model:
            save_dir = 'checkpoint/DANet101_better_model_paddle1.6'
            if os.path.exists(save_dir + '.pdparams'):
                param, _ = fluid.load_dygraph(save_dir)
                model.set_dict(param)
                assert len(param) == len(model.state_dict()), "参数量不一致，加载参数失败，" \
                                                              "请核对模型是否初始化/模型是否一致"
                print('load better model!')

        optimizer = optimizer_setting(args)
        train_data = cityscapes_train(data_root=data_root,
                                      base_size=args.base_size,
                                      crop_size=args.crop_size,
                                      scale=args.scale,
                                      xmap=True,
                                      batch_size=batch_size,
                                      gpu_num=num)
        batch_train_data = paddle.batch(paddle.reader.shuffle(
            train_data, buf_size=batch_size * 64),
            batch_size=batch_size,
            drop_last=True)

        val_data = cityscapes_val(data_root=data_root,
                                  base_size=args.base_size,
                                  crop_size=args.crop_size,
                                  scale=args.scale,
                                  xmap=True)
        batch_test_data = paddle.batch(val_data,
                                       batch_size=batch_size,
                                       drop_last=True)

        train_iou_manager = fluid.metrics.Accuracy()
        train_avg_loss_manager = fluid.metrics.Accuracy()
        test_iou_manager = fluid.metrics.Accuracy()
        test_avg_loss_manager = fluid.metrics.Accuracy()

        better_miou_train = 0
        better_miou_test = 0

        for epoch in range(num_epochs):
            prev_time = datetime.now()
            train_avg_loss_manager.reset()
            train_iou_manager.reset()
            for batch_id, data in enumerate(batch_train_data()):
                image = np.array([x[0] for x in data]).astype('float32')
                label = np.array([x[1] for x in data]).astype('int64')

                image = fluid.dygraph.to_variable(image)
                label = fluid.dygraph.to_variable(label)
                label.stop_gradient = True
                pred, pred2, pred3 = model(image)
                train_loss = loss_fn(pred, pred2, pred3, label, num_classes=num_classes)
                train_avg_loss = fluid.layers.mean(train_loss)
                miou, wrong, correct = mean_iou(pred, label, num_classes=num_classes)
                train_avg_loss.backward()
                optimizer.minimize(train_avg_loss)
                model.clear_gradients()
                train_iou_manager.update(miou.numpy(), weight=batch_size*num)
                train_avg_loss_manager.update(train_avg_loss.numpy(), weight=batch_size*num)
                batch_train_str = "epoch: {}, batch: {}, train_avg_loss: {:.6f}, " \
                                  "train_miou: {:.6f}.".format(epoch + 1,
                                                               batch_id + 1,
                                                               train_avg_loss.numpy()[0],
                                                               miou.numpy()[0])
                if batch_id % 100 == 0:
                    logging.info(batch_train_str)
                    print(batch_train_str)
            cur_time = datetime.now()
            h, remainder = divmod((cur_time - prev_time).seconds, 3600)
            m, s = divmod(remainder, 60)
            time_str = " Time %02d:%02d:%02d" % (h, m, s)
            train_str = "\nepoch: {}, train_avg_loss: {:.6f}, " \
                        "train_miou: {:.6f}.".format(epoch + 1,
                                                     train_avg_loss_manager.eval()[0],
                                                     train_iou_manager.eval()[0])
            print(train_str + time_str + '\n')
            logging.info(train_str + time_str + '\n')
            plot_loss.append(train_loss_title, epoch, train_avg_loss_manager.eval()[0])
            plot_loss.plot('./DANet_loss.jpg')
            plot_iou.append(train_iou_title, epoch, train_iou_manager.eval()[0])
            plot_iou.plot('./DANet_miou.jpg')
            fluid.dygraph.save_dygraph(model.state_dict(), 'checkpoint/DANet_epoch_new')
            # save_model
            if better_miou_train < train_iou_manager.eval()[0]:
                shutil.rmtree('checkpoint/DAnet_better_train_{:.4f}.pdparams'.format(better_miou_train), ignore_errors=True)
                better_miou_train = train_iou_manager.eval()[0]
                fluid.dygraph.save_dygraph(model.state_dict(),
                                           'checkpoint/DAnet_better_train_{:.4f}'.format(better_miou_train))

            ########## test ############
            model.eval()
            test_iou_manager.reset()
            test_avg_loss_manager.reset()
            prev_time = datetime.now()
            for (batch_id, data) in enumerate(batch_test_data()):
                image = np.array([x[0] for x in data]).astype('float32')
                label = np.array([x[1] for x in data]).astype('int64')

                image = fluid.dygraph.to_variable(image)
                label = fluid.dygraph.to_variable(label)

                label.stop_gradient = True
                pred, pred2, pred3 = model(image)
                test_loss = loss_fn(pred, pred2, pred3, label, num_classes=num_classes)
                test_avg_loss = fluid.layers.mean(test_loss)
                miou, wrong, correct = mean_iou(pred, label, num_classes=num_classes)
                test_iou_manager.update(miou.numpy(), weight=batch_size*num)
                test_avg_loss_manager.update(test_avg_loss.numpy(), weight=batch_size*num)
                batch_test_str = "epoch: {}, batch: {}, test_avg_loss: {:.6f}, " \
                                 "test_miou: {:.6f}.".format(epoch + 1, batch_id + 1,
                                                             test_avg_loss.numpy()[0],
                                                             miou.numpy()[0])
                if batch_id % 20 == 0:
                    logging.info(batch_test_str)
                    print(batch_test_str)
            cur_time = datetime.now()
            h, remainder = divmod((cur_time - prev_time).seconds, 3600)
            m, s = divmod(remainder, 60)
            time_str = " Time %02d:%02d:%02d" % (h, m, s)
            test_str = "\nepoch: {}, test_avg_loss: {:.6f}, " \
                       "test_miou: {:.6f}.".format(epoch + 1,
                                                   test_avg_loss_manager.eval()[0],
                                                   test_iou_manager.eval()[0])
            print(test_str + time_str + '\n')
            logging.info(test_str + time_str + '\n')
            plot_loss.append(test_loss_title, epoch, test_avg_loss_manager.eval()[0])
            plot_loss.plot('./DANet_loss.jpg')
            plot_iou.append(test_iou_title, epoch, test_iou_manager.eval()[0])
            plot_iou.plot('./DANet_miou.jpg')
            model.train()
            # save_model
            if better_miou_test < test_iou_manager.eval()[0]:
                shutil.rmtree('checkpoint/DAnet_better_test_{:.4f}.pdparams'.format(better_miou_test), ignore_errors=True)
                better_miou_test = test_iou_manager.eval()[0]
                fluid.dygraph.save_dygraph(model.state_dict(),
                                           'checkpoint/DAnet_better_test_{:.4f}'.format(better_miou_test))


if __name__ == '__main__':
    options = Options()
    args = options.parse()
    options.print_args()
    main(args)