Function main

deploy/slim/quantization/quant.py:89–221

(config, device, logger, vdl_writer)

def main(config, device, logger, vdl_writer):
    # init dist environment
    if config["Global"]["distributed"]:
        dist.init_parallel_env()

    global_config = config["Global"]

    # build dataloader
    set_signal_handlers()
    train_dataloader = build_dataloader(config, "Train", device, logger)
    if config["Eval"]:
        valid_dataloader = build_dataloader(config, "Eval", device, logger)
    else:
        valid_dataloader = None

    # build post process
    post_process_class = build_post_process(config["PostProcess"], global_config)

    # build model
    # for rec algorithm
    if hasattr(post_process_class, "character"):
        char_num = len(getattr(post_process_class, "character"))
        if config["Architecture"]["algorithm"] in [
            "Distillation",
        ]:  # distillation model
            for key in config["Architecture"]["Models"]:
                if (
                    config["Architecture"]["Models"][key]["Head"]["name"] == "MultiHead"
                ):  # for multi head
                    if config["PostProcess"]["name"] == "DistillationSARLabelDecode":
                        char_num = char_num - 2
                    # update SARLoss params
                    assert (
                        list(config["Loss"]["loss_config_list"][-1].keys())[0]
                        == "DistillationSARLoss"
                    )
                    config["Loss"]["loss_config_list"][-1]["DistillationSARLoss"][
                        "ignore_index"
                    ] = (char_num + 1)
                    out_channels_list = {}
                    out_channels_list["CTCLabelDecode"] = char_num
                    out_channels_list["SARLabelDecode"] = char_num + 2
                    config["Architecture"]["Models"][key]["Head"][
                        "out_channels_list"
                    ] = out_channels_list
                else:
                    config["Architecture"]["Models"][key]["Head"][
                        "out_channels"
                    ] = char_num
        elif config["Architecture"]["Head"]["name"] == "MultiHead":  # for multi head
            if config["PostProcess"]["name"] == "SARLabelDecode":
                char_num = char_num - 2
            # update SARLoss params
            assert list(config["Loss"]["loss_config_list"][1].keys())[0] == "SARLoss"
            if config["Loss"]["loss_config_list"][1]["SARLoss"] is None:
                config["Loss"]["loss_config_list"][1]["SARLoss"] = {
                    "ignore_index": char_num + 1
                }