ESPnet源码解析(二)asr_train.py

这部分的代码是声学模型训练的代码的第一部分，以前看的，我代码能力比较薄弱吧，反正只能慢慢改了。
#!/usr/bin/env python3
# encoding: utf-8

# Copyright 2017 Tomoki Hayashi (Nagoya University)
#  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)

"""Automatic speech recognition model training script."""

import logging
import os
import random
import subprocess
import sys

import configargparse
import numpy as np

from espnet import __version__
from espnet.utils.cli_utils import strtobool
from espnet.utils.training.batchfy import BATCH_COUNT_CHOICES

# NOTE: you need this func to generate our sphinx doc
def get_parser(parser=None, required=True):
    """Get default arguments."""
    if parser is None:
        parser = configargparse.ArgumentParser(
            description="Train an automatic speech recognition (ASR) model on one CPU, "
            "one or multiple GPUs",
            config_file_parser_class=configargparse.YAMLConfigFileParser,
            formatter_class=configargparse.ArgumentDefaultsHelpFormatter,
        )
    # general configuration
    parser.add("--config", is_config_file=True, help="config file path")
#(config配置文件)
    parser.add(
        "--config2",
        is_config_file=True,
        help="second config file path that overwrites the settings in `--config`.",
    )
    parser.add(
        "--config3",
        is_config_file=True,
        help="third config file path that overwrites the settings in "
        "`--config` and `--config2`.",
    )

    parser.add_argument(
        "--ngpu",
        default=None,
        type=int,
        help="Number of GPUs. If not given, use all visible devices",
    )
#gpu使用数量
    parser.add_argument(
        "--train-dtype",
        default="float32",
        choices=["float16", "float32", "float64", "O0", "O1", "O2", "O3"],
        help="Data type for training (only pytorch backend). "
        "O0,O1,.. flags require apex. "
        "See https://nvidia.github.io/apex/amp.html#opt-levels",
    )
#训练数据类型
    parser.add_argument(
        "--backend",
        default="chainer",
        type=str,
        choices=["chainer", "pytorch"],
        help="Backend library",
    )
#指定训练框架，pytorch或者chainer
    parser.add_argument(
        "--outdir", type=str, required=required, help="Output directory"
    )
#输出目录的文件夹
    parser.add_argument("--debugmode", default=1, type=int, help="Debugmode")
#debugmode的数量
    parser.add_argument("--dict", required=required, help="Dictionary")
#字典
    parser.add_argument("--seed", default=1, type=int, help="Random seed")
#随机速度（不知道什么意思）
    parser.add_argument("--debugdir", type=str, help="Output directory for debugging")
#debug的输出目录
    parser.add_argument(
        "--resume",
        "-r",
        default="",
        nargs="?",
        help="Resume the training from snapshot",
    )
#从指定的已训练好的模型继续训练
    parser.add_argument(
        "--minibatches",
        "-N",
        type=int,
        default="-1",
        help="Process only N minibatches (for debug)",
    )
#小批量训练，详细可以百度
    parser.add_argument("--verbose", "-V", default=0, type=int, help="Verbose option")
#日志选项
    parser.add_argument(
        "--tensorboard-dir",
        default=None,
        type=str,
        nargs="?",
        help="Tensorboard log dir path",
    )
#Tensorboard日志目录存放路径，可以学一学
    parser.add_argument(
        "--report-interval-iters",
        default=100,
        type=int,
        help="Report interval iterations",
    )
#迭代多少次输出一次，默认100 
    parser.add_argument(
        "--save-interval-iters",
        default=0,
        type=int,
        help="Save snapshot interval iterations",
    )
#保存模型训练的初始迭代
    # task related
    parser.add_argument(
        "--train-json",
        type=str,
        default=None,
        help="Filename of train label data (json)",
    )
#训练数据类型-jaon
    parser.add_argument(
        "--valid-json",
        type=str,
        default=None,
        help="Filename of validation label data (json)",
    )
#验证集数据类型-jaon
    # network architecture
    parser.add_argument(
        "--model-module",
        type=str,
        default=None,
        help="model defined module (default: espnet.nets.xxx_backend.e2e_asr:E2E)",
    )
#模型定义模块，用pytorch的还是chainer的
    # encoder
    parser.add_argument(
        "--num-encs", default=1, type=int, help="Number of encoders in the model."
    )
#模型的编码器数量
    # loss related
    parser.add_argument(
        "--ctc_type",
        default="warpctc",
        type=str,
        choices=["builtin", "warpctc", "gtnctc", "cudnnctc"],
        help="Type of CTC implementation to calculate loss.",
    )
#计算CTC的损失是用什么模型实现的，我用的是warpctc
    parser.add_argument(
        "--mtlalpha",
        default=0.5,
        type=float,
        help="Multitask learning coefficient, "
        "alpha: alpha*ctc_loss + (1-alpha)*att_loss ",
    )
#多任务学习系数，公式为上面那个
    parser.add_argument(
        "--lsm-weight", default=0.0, type=float, help="Label smoothing weight"
    )
#标签平滑权重
    # recognition options to compute CER/WER
    parser.add_argument(
        "--report-cer",
        default=False,
        action="store_true",
        help="Compute CER on development set",
    )
#计算验证集的字错误
    parser.add_argument(
        "--report-wer",
        default=False,
        action="store_true",
        help="Compute WER on development set",
    )
#计算验证集的词错误率
    parser.add_argument("--nbest", type=int, default=1, help="Output N-best hypotheses")
#输出几个最好的假设，默认为1 
    parser.add_argument("--beam-size", type=int, default=4, help="Beam size")
#beam search默认设为4
    parser.add_argument("--penalty", default=0.0, type=float, help="Incertion penalty")
#插入惩罚参数，默认为0
    parser.add_argument(
        "--maxlenratio",
        default=0.0,
        type=float,
        help="""Input length ratio to obtain max output length.
                        If maxlenratio=0.0 (default), it uses a end-detect function
                        to automatically find maximum hypothesis lengths""",
#输出目标句子最长和句子的最大比，详细看论文混合CTC/attention
    )
    parser.add_argument(
        "--minlenratio",
        default=0.0,
        type=float,
        help="Input length ratio to obtain min output length",
    )
#同上
    parser.add_argument(
        "--ctc-weight", default=0.3, type=float, help="CTC weight in joint decoding"
    )
#CTC在联合解码中的权重占比，默认为0.3
    parser.add_argument(
        "--rnnlm", type=str, default=None, help="RNNLM model file to read"
    )
#RNN语言模型要读取的文件
    parser.add_argument(
        "--rnnlm-conf", type=str, default=None, help="RNNLM model config file to read"
    )
#RNN语言模型默认配置
    parser.add_argument("--lm-weight", default=0.1, type=float, help="RNNLM weight.")
#RNN语言模型在解码中的占比权重
    parser.add_argument("--sym-space", default="<space>", type=str, help="Space symbol")
#空格符号用<space>代替
    parser.add_argument("--sym-blank", default="<blank>", type=str, help="Blank symbol")
#空白符号用<blank>符号代替
    # minibatch related
    parser.add_argument(
        "--sortagrad",
        default=0,
        type=int,
        nargs="?",
        help="How many epochs to use sortagrad for. 0 = deactivated, -1 = all epochs",
    )
#minibatch相关，多少次epochs遍历完一次
    parser.add_argument(
        "--batch-count",
        default="auto",
        choices=BATCH_COUNT_CHOICES,
        help="How to count batch_size. "
        "The default (auto) will find how to count by args.",
    )
#batchsize大小
    parser.add_argument(
        "--batch-size",
        "--batch-seqs",
        "-b",
        default=0,
        type=int,
        help="Maximum seqs in a minibatch (0 to disable)",
    )
#
    parser.add_argument(
        "--batch-bins",
        default=0,
        type=int,
        help="Maximum bins in a minibatch (0 to disable)",
    )
    parser.add_argument(
        "--batch-frames-in",
        default=0,
        type=int,
        help="Maximum input frames in a minibatch (0 to disable)",
    )
    parser.add_argument(
        "--batch-frames-out",
        default=0,
        type=int,
        help="Maximum output frames in a minibatch (0 to disable)",
    )
    parser.add_argument(
        "--batch-frames-inout",
        default=0,
        type=int,
        help="Maximum input+output frames in a minibatch (0 to disable)",
    )
    parser.add_argument(
        "--maxlen-in",
        "--batch-seq-maxlen-in",
        default=800,
        type=int,
        metavar="ML",
        help="When --batch-count=seq, "
        "batch size is reduced if the input sequence length > ML.",
    )
    parser.add_argument(
        "--maxlen-out",
        "--batch-seq-maxlen-out",
        default=150,
        type=int,
        metavar="ML",
        help="When --batch-count=seq, "
        "batch size is reduced if the output sequence length > ML",
    )
    parser.add_argument(
        "--n-iter-processes",
        default=0,
        type=int,
        help="Number of processes of iterator",
    )
    parser.add_argument(
        "--preprocess-conf",
        type=str,
        default=None,
        nargs="?",
        help="The configuration file for the pre-processing",
    )
    # optimization related
    parser.add_argument(
        "--opt",
        default="adadelta",
        type=str,
        choices=["adadelta", "adam", "noam"],
        help="Optimizer",
    )
#优化模型，"adadelta", "adam", "noam"三个选项，默认adadelta
    parser.add_argument(
        "--accum-grad", default=1, type=int, help="Number of gradient accumuration"
    )
#梯度累计次数
    parser.add_argument(
        "--eps", default=1e-8, type=float, help="Epsilon constant for optimizer"
    )
#优化器的epsilon系数，因为有的鬼地方防止除0，比如BN
    parser.add_argument(
        "--eps-decay", default=0.01, type=float, help="Decaying ratio of epsilon"
    )
#优化器epsilon衰减比率
    parser.add_argument(
        "--weight-decay", default=0.0, type=float, help="Weight decay ratio"
    )
#权重衰减比率
    parser.add_argument(
        "--criterion",
        default="acc",
        type=str,
        choices=["loss", "loss_eps_decay_only", "acc"],
        help="Criterion to perform epsilon decay",
    )
#标准epsilon衰减
    parser.add_argument(
        "--threshold", default=1e-4, type=float, help="Threshold to stop iteration"
    )
#停止迭代的阈值
    parser.add_argument(
        "--epochs", "-e", default=30, type=int, help="Maximum number of epochs"
    )
#epochs次数
    parser.add_argument(
        "--early-stop-criterion",
        default="validation/main/acc",
        type=str,
        nargs="?",
        help="Value to monitor to trigger an early stopping of the training",
    )
#监控触发停止训练的值
    parser.add_argument(
        "--patience",
        default=3,
        type=int,
        nargs="?",
        help="Number of epochs to wait without improvement "
        "before stopping the training",
    )
#没有再优化模型的epochs的数量，然后提前结束训练
    parser.add_argument(
        "--grad-clip", default=5, type=float, help="Gradient norm threshold to clip"
    )

    parser.add_argument(
        "--num-save-attention",
        default=3,
        type=int,
        help="Number of samples of attention to be saved",
    )
#保留注意力样本的数量，可以看result
    parser.add_argument(
        "--num-save-ctc",
        default=3,
        type=int,
        help="Number of samples of CTC probability to be saved",
    )
#要保留ctc概率的数量，也可以看result
    parser.add_argument(
        "--grad-noise",
        type=strtobool,
        default=False,
        help="The flag to switch to use noise injection to gradients during training",
    )
#梯度时加噪
    # asr_mix related
    parser.add_argument(
        "--num-spkrs",
        default=1,
        type=int,
        choices=[1, 2],
        help="Number of speakers in the speech.",
    )
#语音中说话人的数量
    # decoder related
    parser.add_argument(
        "--context-residual",
        default=False,
        type=strtobool,
        nargs="?",
        help="The flag to switch to use context vector residual in the decoder network",
    )
#使用上下文残差
    # finetuning related
    parser.add_argument(
        "--enc-init",
        default=None,
        type=str,
        help="Pre-trained ASR model to initialize encoder.",
    )
#预训练语音识别模型的初始化编码器
    parser.add_argument(
        "--enc-init-mods",
        default="enc.enc.",
        type=lambda s: [str(mod) for mod in s.split(",") if s != ""],
        help="List of encoder modules to initialize, separated by a comma.",
    )
#要初始化的编码器模块
    parser.add_argument(
        "--dec-init",
        default=None,
        type=str,
        help="Pre-trained ASR, MT or LM model to initialize decoder.",
    )
#预训练语音识别机器翻译和语言模型初始化编码器
    parser.add_argument(
        "--dec-init-mods",
        default="att.,dec.",
        type=lambda s: [str(mod) for mod in s.split(",") if s != ""],
        help="List of decoder modules to initialize, separated by a comma.",
    )
#初始化编码器模块
    parser.add_argument(
        "--freeze-mods",
        default=None,
        type=lambda s: [str(mod) for mod in s.split(",") if s != ""],
        help="List of modules to freeze, separated by a comma.",
    )
    # front end related
    parser.add_argument(
        "--use-frontend",
        type=strtobool,
        default=False,
        help="The flag to switch to use frontend system.",
    )
#这个标志意味着使用前端系统
    # WPE related
    parser.add_argument(
        "--use-wpe",
        type=strtobool,
        default=False,
        help="Apply Weighted Prediction Error",
    )
#应用权重预测误差，作用是去混响
    parser.add_argument(
        "--wtype",
        default="blstmp",
        type=str,
        choices=[
            "lstm",
            "blstm",
            "lstmp",
            "blstmp",
            "vgglstmp",
            "vggblstmp",
            "vgglstm",
            "vggblstm",
            "gru",
            "bgru",
            "grup",
            "bgrup",
            "vgggrup",
            "vggbgrup",
            "vgggru",
            "vggbgru",
        ],
        help="Type of encoder network architecture "
        "of the mask estimator for WPE. "
        "",
    )
#编码网络类别类别
    parser.add_argument("--wlayers", type=int, default=2, help="")#层数
    parser.add_argument("--wunits", type=int, default=300, help="")#神经元个数
    parser.add_argument("--wprojs", type=int, default=300, help="")
    parser.add_argument("--wdropout-rate", type=float, default=0.0, help="")
    parser.add_argument("--wpe-taps", type=int, default=5, help="")
    parser.add_argument("--wpe-delay", type=int, default=3, help="")
    parser.add_argument(
        "--use-dnn-mask-for-wpe",
        type=strtobool,
        default=False,
        help="Use DNN to estimate the power spectrogram. "
        "This option is experimental.",
    )
    # Beamformer related
    parser.add_argument("--use-beamformer", type=strtobool, default=True, help="")
    parser.add_argument(
        "--btype",
        default="blstmp",
        type=str,
        choices=[
            "lstm",
            "blstm",
            "lstmp",
            "blstmp",
            "vgglstmp",
            "vggblstmp",
            "vgglstm",
            "vggblstm",
            "gru",
            "bgru",
            "grup",
            "bgrup",
            "vgggrup",
            "vggbgrup",
            "vgggru",
            "vggbgru",
        ],
        help="Type of encoder network architecture "
        "of the mask estimator for Beamformer.",
    )
    parser.add_argument("--blayers", type=int, default=2, help="")
    parser.add_argument("--bunits", type=int, default=300, help="")
    parser.add_argument("--bprojs", type=int, default=300, help="")
    parser.add_argument("--badim", type=int, default=320, help="")
    parser.add_argument(
        "--bnmask",
        type=int,
        default=2,
        help="Number of beamforming masks, " "default is 2 for [speech, noise].",
    )
    parser.add_argument(
        "--ref-channel",
        type=int,
        default=-1,
        help="The reference channel used for beamformer. "
        "By default, the channel is estimated by DNN.",
    )
    parser.add_argument("--bdropout-rate", type=float, default=0.0, help="")
    # Feature transform: Normalization
    parser.add_argument(
        "--stats-file",
        type=str,
        default=None,
        help="The stats file for the feature normalization",
    )
    parser.add_argument(
        "--apply-uttmvn",
        type=strtobool,
        default=True,
        help="Apply utterance level mean " "variance normalization.",
    )
    parser.add_argument("--uttmvn-norm-means", type=strtobool, default=True, help="")
    parser.add_argument("--uttmvn-norm-vars", type=strtobool, default=False, help="")
    # Feature transform: Fbank
    parser.add_argument(
        "--fbank-fs",
        type=int,
        default=16000,
        help="The sample frequency used for " "the mel-fbank creation.",
    )
    parser.add_argument(
        "--n-mels", type=int, default=80, help="The number of mel-frequency bins."
    )
    parser.add_argument("--fbank-fmin", type=float, default=0.0, help="")
    parser.add_argument("--fbank-fmax", type=float, default=None, help="")
    return parser

def main(cmd_args):
    """Run the main training function."""
    parser = get_parser()
#获取参数
    args, _ = parser.parse_known_args(cmd_args)
#多次传参
    if args.backend == "chainer" and args.train_dtype != "float32":
        raise NotImplementedError(
            f"chainer backend does not support --train-dtype {args.train_dtype}."
            "Use --dtype float32."
        )
#如果选择chainer框架且训练类型不是float32则报错
    if args.ngpu == 0 and args.train_dtype in ("O0", "O1", "O2", "O3", "float16"):
        raise ValueError(
            f"--train-dtype {args.train_dtype} does not support the CPU backend."
        )
#如果使用CPU训练且训练类型为O0-O3等报错
    from espnet.utils.dynamic_import import dynamic_import
#从espnet_utils_dynamic_import.py下引用方法dynamic_import，作用动态引入模块
    if args.model_module is None:
        if args.num_spkrs == 1:
            model_module = "espnet.nets." + args.backend + "_backend.e2e_asr:E2E"
        else:
            model_module = "espnet.nets." + args.backend + "_backend.e2e_asr_mix:E2E"
    else:
        model_module = args.model_module
    model_class = dynamic_import(model_module)
#代码注释见https://shimo.im/docs/9030MPWRzYUNZrqw,这里是指e2e_asr_transformer.py中的E2E类
    model_class.add_arguments(parser)
#将parser参数传给模型，详情见https://shimo.im/docs/vVAXVoDOYNFVVjqm

    args = parser.parse_args(cmd_args)
    args.model_module = model_module
    if "chainer_backend" in args.model_module:
        args.backend = "chainer"
    if "pytorch_backend" in args.model_module:
        args.backend = "pytorch"
#指定框架,这里用的是pytorch

    # add version info in args
    args.version = __version__
#添加版本信息
    # logging info
    if args.verbose > 0:
        logging.basicConfig(
            level=logging.INFO,
            format="%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s",
        )
    else:
        logging.basicConfig(
            level=logging.WARN,
            format="%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s",
        )
        logging.warning("Skip DEBUG/INFO messages")

    # If --ngpu is not given,
    #   1. if CUDA_VISIBLE_DEVICES is set, all visible devices
    #   2. if nvidia-smi exists, use all devices
    #   3. else ngpu=0
    if args.ngpu is None:
        cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
#获取能用的GPU
        if cvd is not None:
            ngpu = len(cvd.split(","))
#如果cvd不为空，按“，”切分，获取gpu数量
        else:
            logging.warning("CUDA_VISIBLE_DEVICES is not set.")
            try:
                p = subprocess.run(
                    ["nvidia-smi", "-L"], stdout=subprocess.PIPE, stderr=subprocess.PIPE
                )
            except (subprocess.CalledProcessError, FileNotFoundError):
                ngpu = 0
#警告GPU没设置
            else:
                ngpu = len(p.stderr.decode().split("\n")) - 1
    else:
        if args.ngpu != 1:
            logging.debug(
                "There are some bugs with multi-GPU processing in PyTorch 1.2+"
                + " (see https://github.com/pytorch/pytorch/issues/21108)"
            )
        ngpu = args.ngpu
    logging.info(f"ngpu: {ngpu}")
#如果GPU数量不等于1，提出警告需要pytorch1.2+,
   # display PYTHONPATH
    logging.info("python path = " + os.environ.get("PYTHONPATH", "(None)"))
#python路径
    # set random seed
    logging.info("random seed = %d" % args.seed)
    random.seed(args.seed)
    np.random.seed(args.seed)
#设置随机seed，就是避免二次调用的时候产生不同的随机数据集。你再问细一点我也不知道
    # load dictionary for debug log
    if args.dict is not None:
        with open(args.dict, "rb") as f:
            dictionary = f.readlines()
#如果字典不为空，按行读取一行就是长这样“一 2”，
        char_list = [entry.decode("utf-8").split(" ")[0] for entry in dictionary]
        #前面字符后面数字映射，按空格切分取出字  
        char_list.insert(0, "<blank>")
        #在索引为0的位置插入<blank>,就是第一个位置插入<blank>
        char_list.append("<eos>")
        #在最后的位置插入<eos>
        # for non-autoregressive maskctc model
        if "maskctc" in args.model_module:
            char_list.append("<mask>")
        #参考论文：Mask CTC: Non-Autoregressive End-to-End ASR with CTC and Mask Predict
        args.char_list = char_list
        #重新将字典赋值给args
    else:
        args.char_list = None

    # train
    logging.info("backend = " + args.backend)

    if args.num_spkrs == 1:
        if args.backend == "chainer":
            from espnet.asr.chainer_backend.asr import train

            train(args)
        elif args.backend == "pytorch":
            from espnet.asr.pytorch_backend.asr import train

            train(args)
        else:
            raise ValueError("Only chainer and pytorch are supported.")
    else:
        # FIXME(kamo): Support --model-module
        if args.backend == "pytorch":
            from espnet.asr.pytorch_backend.asr_mix import train

            train(args)
        else:
            raise ValueError("Only pytorch is supported.")
#训练,详情见https://shimo.im/docs/5xkGMLnEE9cQxp3X
if __name__ == "__main__":
    main(sys.argv[1:])
ESPnet
ESPnet ESPnet_train
本博客所有文章除特别声明外，均采用 CC BY-SA 4.0 协议，转载请注明出处！
端到端语音识别模型中的LM Fusion 上一篇
ESPnet源码解析(一)run.sh 下一篇