Function build_model

CV/SemSegPaddle/src/models/model_builder.py:113–255

(main_prog, start_prog, phase=ModelPhase.TRAIN)

def build_model(main_prog, start_prog, phase=ModelPhase.TRAIN):
    if not ModelPhase.is_valid_phase(phase):
        raise ValueError("ModelPhase {} is not valid!".format(phase))
    if ModelPhase.is_train(phase):
        width = cfg.DATAAUG.CROP_SIZE
        height = cfg.DATAAUG.CROP_SIZE
    else:
        width = cfg.TEST.CROP_SIZE
        height = cfg.TEST.CROP_SIZE

    image_shape = [cfg.DATASET.DATA_DIM, height, width]
    grt_shape = [1, height, width]
    class_num = cfg.DATASET.NUM_CLASSES

    with fluid.program_guard(main_prog, start_prog):
        with fluid.unique_name.guard():
            # 在导出模型的时候，增加图像标准化预处理,减小预测部署时图像的处理流程
            # 预测部署时只须对输入图像增加batch_size维度即可
            if ModelPhase.is_predict(phase):
                origin_image = fluid.layers.data(name='image', 
                        shape=[ -1, 1, 1, cfg.DATASET.DATA_DIM], 
                        dtype='float32', 
                        append_batch_size=False)
                image = fluid.layers.transpose(origin_image, [0, 3, 1, 2])
                origin_shape = fluid.layers.shape(image)[-2:]
                mean = np.array(cfg.MEAN).reshape(1, len(cfg.MEAN), 1, 1)
                mean = fluid.layers.assign(mean.astype('float32'))
                std = np.array(cfg.STD).reshape(1, len(cfg.STD), 1, 1)
                std = fluid.layers.assign(std.astype('float32'))
                image = (image/255 - mean)/std
                image = fluid.layers.resize_bilinear(image, 
                        out_shape=[height, width], align_corners=False, align_mode=0)
            else:
                image = fluid.layers.data( name='image', shape=image_shape, dtype='float32')
            label = fluid.layers.data( name='label', shape=grt_shape, dtype='int32')
            mask = fluid.layers.data( name='mask', shape=grt_shape, dtype='int32')

            # use PyReader when doing traning and evaluation
            if ModelPhase.is_train(phase) or ModelPhase.is_eval(phase):
                iterable = True if ModelPhase.is_eval(phase) else False
                print("iterable: ", iterable)
                py_reader = fluid.io.PyReader(
                    feed_list=[image, label, mask],
                    capacity=cfg.DATALOADER.BUF_SIZE,
                    iterable=iterable,
                    use_double_buffer=True,
                    return_list=False)

            model_name = map_model_name(cfg.MODEL.MODEL_NAME)
            model_func = get_func("modeling." + model_name)

            loss_type = cfg.SOLVER.LOSS
            if not isinstance(loss_type, list):
                loss_type = list(loss_type)

            # dice_loss或bce_loss只适用两类分割中
            if class_num > 2 and (("dice_loss" in loss_type) or ("bce_loss" in loss_type)):
                raise Exception("dice loss and bce loss is only applicable to binary classfication")