Pytorch
TensorFlow

#
从头开始训练因果语言模型

到目前为止，我们主要使用预训练模型，并通过复用预训练的权重，然后使用新的数据对它们进行微调，以适应新的应用场景。正如我们在 第一章 中看到的，这通常称为 `迁移学习（transfer learning）` ，对于大多数标注数据稀缺的应用场景，它是一种将 Transformer 模型应用到大部分真实的应用场景中的一个非常成功的策略。在本章中，我们将采用不同的方法并从头开始训练一个全新的模型。如果你有大量数据而且这些数据与可用模型的预训练数据差异很大，那么这是一个很好的方法。然而，相比仅微调现有模型，预训练语言模型需要更多的计算资源。训练一个新模型可能是有意义的示例包括由音乐符号、DNA 等分子序列或编程语言组成的数据集。编程语言组成的数据集最近广泛地受到关注，这要归功于 TabNine 和 GitHub 的 Copilot 等工具的流行，它们由 OpenAI 的 Codex 模型提供支持，可以生成长代码序列。这种文本生成任务最适合使用自回归或因果语言模型（例如 GPT-2）。

在这一节，我们将构建一个精简版的代码生成模型：使用 Python 代码的一个数据集，来实现一行代码的补全，而不是直接生成完整的函数或类。当你使用 Python 处理数据时，你经常会接触到 Python 数据科学栈，包括 `matplotlib` ， `seaborn` ， `pandas` ，和 `scikit-learn` 这些库。当使用这些框架时，经常需要查找特定的命令，如果我们能够用模型来自动给出恰当的推荐命令就太好了！

在 第六章 中，我们创建了一个高效的 tokenizer 来处理 Python 源代码，但我们还需要一个大规模的数据集来预训练模型。在这里，我们将使用 tokenizer 处理一个来自 GitHub 仓库的 Python 代码语料库。然后，我们将使用 Trainer API 和 🤗 Accelerate 来训练模型。让我们开始吧！

这里展示的是一个已经训练并上传到 Hub 的模型，它就是使用本节中的代码训练的。你可以在 这里 找到它。注意，由于文本生成过程中有一些随机性，你可能会得到稍微不同的结果。

##
收集数据

我们可以从诸如 GitHub 这样的代码仓库中获取丰富的 Python 代码，通过对每个 Python 仓库进行抓取，我们就可以创建一个数据集。这就是在 Transformers textbook 中预训练一个大型 GPT-2 模型的方法。开发者整理了名为 `codeparrot` 的一个大约为 180GB 的 GitHub 数据集， 其中包含大约 2,000 万个的Python 文件。 开发者用这些文件构建了一个数据集，并在 Hugging Face Hub 上分享了这个数据集。

然而，使用完整语料库的训练既耗时又费力，我们只需要找到 Python 数据科学栈相关的数据集子集。所以，让我们从 `codeparrot` 数据集中筛选出包含这个栈中所有相关库的所有文件。由于数据集的太大，我们希望避免直接把全部的数据集下载下来；因此，我们将使用流式传输的方法来动态过滤它。为了使用上述的库来筛选代码样本，我们将使用以下函数：

    def any_keyword_in_string(string, keywords):
        for keyword in keywords:
            if keyword in string:
                return True
        return False

让我们用两个例子来测试一下：

    filters = ["pandas", "sklearn", "matplotlib", "seaborn"]
    example_1 = "import numpy as np"
    example_2 = "import pandas as pd"

    print(
        any_keyword_in_string(example_1, filters), any_keyword_in_string(example_2, filters)
    )

    False True

我们可以使用这个函数来创建一个新的函数，该函数将流式传输数据集并过滤我们想要的元素：

    from collections import defaultdict
    from tqdm import tqdm
    from datasets import Dataset

    def filter_streaming_dataset(dataset, filters):
        filtered_dict = defaultdict(list)
        total = 0
        for sample in tqdm(iter(dataset)):
            total += 1
            if any_keyword_in_string(sample["content"], filters):
                for k, v in sample.items():
                    filtered_dict[k].append(v)
        print(f"{len(filtered_dict['content'])/total:.2%} of data after filtering.")
        return Dataset.from_dict(filtered_dict)

然后我们可以直接使用这里函数流式处理数据集：

    # 执行这个代码块需要非常长的时间,因此你可以跳过它,继续执行下一个!
    from datasets import load_dataset

    split = "train"  # "valid"
    filters = ["pandas", "sklearn", "matplotlib", "seaborn"]

    data = load_dataset(f"transformersbook/codeparrot-{split}", split=split, streaming=True)
    filtered_data = filter_streaming_dataset(data, filters)

    3.26% of data after filtering.

完成这个操作后，我们过滤后的数据集只有原始数据集的大约 3％，但这仍然是相当可观的大小——最终的数据集是 6GB，由 600,000 个 Python 脚本组成！

过滤完整的数据集可能需要 2-3 小时，这取决于你的机器性能和带宽。如果你不想亲自经历这个漫长的过程，我们在 Hub 上提供了过滤后的数据集供你下载：

    from datasets import load_dataset, DatasetDict

    ds_train = load_dataset("huggingface-course/codeparrot-ds-train", split="train")
    ds_valid = load_dataset("huggingface-course/codeparrot-ds-valid", split="validation")

    raw_datasets = DatasetDict(
        {
            "train": ds_train,  # .shuffle().select(range(50000)),
            "valid": ds_valid,  # .shuffle().select(range(500))
        }
    )

    raw_datasets

    DatasetDict({
        train: Dataset({
            features: ['repo_name', 'path', 'copies', 'size', 'content', 'license'],
            num_rows: 606720
        })
        valid: Dataset({
            features: ['repo_name', 'path', 'copies', 'size', 'content', 'license'],
            num_rows: 3322
        })
    })

让我们看一个来自数据集的例子。我们将只显示每个字段的前 200 个字符：

    for key in raw_datasets["train"][0]:
        print(f"{key.upper()}: {raw_datasets['train'][0][key][:200]}")

    'REPO_NAME: kmike/scikit-learn'
    'PATH: sklearn/utils/__init__.py'
    'COPIES: 3'
    'SIZE: 10094'
    '''CONTENT: """
    The :mod:`sklearn.utils` module includes various utilites.
    """

    from collections import Sequence

    import numpy as np
    from scipy.sparse import issparse
    import warnings

    from .murmurhash import murm
    LICENSE: bsd-3-clause'''

我们可以看到， `content` 字段包含了我们希望模型训练的代码。有了这个数据集之后，我们需要对文本进行一些处理，以便它们适合于预训练。

##
准备数据集

首先，我们需要将数据进行分词处理，这样才能进行训练。由于我们的主要目标是自动补全短的函数调用，因此我们可以将上下文大小设置得相对较小。这样做的好处是我们可以更快地训练模型，而且需要的内存也大大减少。如果你的应用需要更多的上下文（比如，你希望模型根据包含函数定义的文件编写单元测试），那么应该增大该数字，但是也要记住这会增加 GPU 显存的占用。现在，我们将上下文大小固定为 128 个 tokens 而不是在 GPT-2 或 GPT-3 中使用的 1,024 或 2,048 个 tokens

大多数文档都包含超过 128 个 tokens 因此简单地将输入截断到最大长度会删除我们数据集的很一大部分。因此，我们将使用 `return_overflowing_tokens` 选项将整个输入进行分词处理，并将其分割为几个块，正如我们在 第六章 中所做的那样。我们还将使用 `return_length` 选项自动返回创建的每个块的长度。通常，最后一个块的大小会小于上下文大小，我们将去掉最后一块以避免填充问题；因为我们已经有足够的数据，所以不需要它们。

让我们通过查看前两个示例来具体了解结果怎么样：

    from transformers import AutoTokenizer

    context_length = 128
    tokenizer = AutoTokenizer.from_pretrained("huggingface-course/code-search-net-tokenizer")

    outputs = tokenizer(
        raw_datasets["train"][:2]["content"],
        truncation=True,
        max_length=context_length,
        return_overflowing_tokens=True,
        return_length=True,
    )

    print(f"Input IDs length: {len(outputs['input_ids'])}")
    print(f"Input chunk lengths: {(outputs['length'])}")
    print(f"Chunk mapping: {outputs['overflow_to_sample_mapping']}")

    Input IDs length: 34
    Input chunk lengths: [128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 117, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 41]
    Chunk mapping: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]

我们可以看到，这两个例子总共得到了 34 个块。查看块长度，我们可以看到两个文档末端的块少于 128 个 tokens （分别为 117 和 41）。不过这些只占我们所拥有的总块数的一小部分，因此我们可以放心地丢掉它们。通过 `overflow_to_sample_mapping` 字段，我们还可以分辨出哪些块属于哪个样本。

在这个操作中，我们使用了🤗 Datasets 中的 `Dataset.map()` 函数的一个便捷的特性，即它并不需要一对一地设置分块后和分块前的映射关系；正如我们在 第三节 中看到的，我们可以自由地将一个样本拆分或者删除部分样本来创建比输入的 `batch_size` 更多或更少元素的 batch。 `Dataset.map()` 函数会自动帮我们关联映射关系，当进行像数据增强或数据过滤这样改变元素数量的操作时非常有用。在我们的情况下，当将每个样本分词并分割成指定上下文大小的块时，我们从每个样本中创建了许多样本。我们需要删除原本的列，因为它们的大小和我们分割后的大小不一样。如果我们想保留它们，我们可以复制它们来填充，并在 `Dataset.map()` 调用中返回它们。

    def tokenize(element):
        outputs = tokenizer(
            element["content"],
            truncation=True,
            max_length=context_length,
            return_overflowing_tokens=True,
            return_length=True,
        )
        input_batch = []
        for length, input_ids in zip(outputs["length"], outputs["input_ids"]):
            if length == context_length:
                input_batch.append(input_ids)
        return {"input_ids": input_batch}

    tokenized_datasets = raw_datasets.map(
        tokenize, batched=True, remove_columns=raw_datasets["train"].column_names
    )
    tokenized_datasets

    DatasetDict({
        train: Dataset({
            features: ['input_ids'],
            num_rows: 16702061
        })
        valid: Dataset({
            features: ['input_ids'],
            num_rows: 93164
        })
    })

我们现在有 1670 万个样本，每个样本有 128 个 tokens 总共相当于大约 21 亿个 tokens。作为参考，OpenAI 的 GPT-3 和 Codex 模型分别在 300 和 1000 亿个 tokens 上进行了训练，其中 Codex 模型从 GPT-3 checkpoint 初始化。本节的目标不是与这些能生成长且连贯文本的模型竞争，而是创建一个能为数据科学家提供快速自动代码补全功能的精简版本。

既然我们已经准备好了数据集，那就来设置模型吧！

✏️ 试一试！这里我们删除了所有小于设定的上下文大小的块，并不会造成大问题，因为我们使用的是比较小的上下文窗口。随着增大上下文大小（或者语料库中的文档长度都很短），被抛弃的块的比例也会增加。更有效方法是将所有 tokenize 后的样本拼接起来加入一个 batch 中，每个样本之间有一个 `eos_token_id` token 作为分隔，然后对连接后的序列进行切块处理。作为练习，修改 `tokenize()` 函数以利用这种方法。请注意，为了获取完整的 token ID 序列你需要设置 `truncation=False` ，并删除 tokenizer 中的其他参数。

##
初始化一个新模型

我们的第一步是初始化一个全新地 GPT-2 模型。我们可以通过加载预训练配置来初始化一个与 GPT-2 small 相同的配置的模型，并确保 tokenizer 大小与模型的词汇表大小匹配，以及设置 `bos` 和 `eos` （序列的开始和结束） token IDs：

    from transformers import AutoTokenizer, GPT2LMHeadModel, AutoConfig

    config = AutoConfig.from_pretrained(
        "gpt2",
        vocab_size=len(tokenizer),
        n_ctx=context_length,
        bos_token_id=tokenizer.bos_token_id,
        eos_token_id=tokenizer.eos_token_id,
    )

有了这个配置对象，我们就可以加载一个全新的 GPT-2 模型。注意，这是我们第一次不使用 `from_pretrained()` 函数，因为我们实际上是自己初始化一个全新的模型而不是从一个预训练的模型继续训练：

    model = GPT2LMHeadModel(config)
    model_size = sum(t.numel() for t in model.parameters())
    print(f"GPT-2 size: {model_size/1000**2:.1f}M parameters")

    GPT-2 size: 124.2M parameters

我们的新模型有 124M 个参数需要训练。在开始训练之前，我们需要设置一个数据整理器（DataCollator），它将负责创建 Batch。我们可以使用 `DataCollatorForLanguageModeling` ，顾名思义，它专门用于语言建模。除了堆叠和填充创建 Batch 之外，它还负责创建语言模型的待预测的标签 —— 在因果语言建模中，输入就是待预测的标签（只是偏移一个元素），而这个数据整理器（DataCollator）在训练过程中实时将输入偏移一个元素来创建它们，因此我们不需要复制 `input_ids` 。

注意， `DataCollatorForLanguageModeling` 同时支持掩码语言建模 （MLM） 和因果语言建模 （CLM）。默认情况下它安装 MLM 需要的格式准备数据，但我们可以通过设置 `mlm=False` 参数切换到 CLM。

    from transformers import DataCollatorForLanguageModeling

    tokenizer.pad_token = tokenizer.eos_token
    data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)

让我们看一个例子：

    out = data_collator([tokenized_datasets["train"][i] for i in range(5)])
    for key in out:
        print(f"{key} shape: {out[key].shape}")

    input_ids shape: torch.Size([5, 128])
    attention_mask shape: torch.Size([5, 128])
    labels shape: torch.Size([5, 128])

我们可以看到示例的数据已经处理好了，并且所有 tensor 都具有相同的形状。

⚠️ 输入序列和目标序列对齐将在模型内部自动进行，所以数据整理器只需复制输入序列来创建目标序列。

现在我们已经准备好了所有东西，可以开始训练我们的模型了——好像也不是那么困难！在我们开始训练之前，我们应该登录到 Hugging Face。如果你正在使用 Notebook 运行代码，你可以使用下面的实用函数进行登录：

    from huggingface_hub import notebook_login

    notebook_login()

这将显示一个小部件，你可以在其中输入你的 Hugging Face 登录凭据。

如果你不是在 Notebook 上工作，只需在终端中输入以下行：

    huggingface-cli login

剩下要做的就是配置训练参数并启动 `Trainer` 。本次的训练中我们将使用余弦学习率调度，并进行一些 Warmup。训练的 batch size 是 256 （ `per_device_train_batch_size` * `gradient_accumulation_steps` ）。当单个 batch 无法放入内存时，可以使用梯度累积，并通过多次向前/向后传递逐步累积梯度。当我们在本节最后使用 🤗 Accelerate 创建训练循环时，我们将看到这一点。

    from transformers import Trainer, TrainingArguments

    args = TrainingArguments(
        output_dir="codeparrot-ds",
        per_device_train_batch_size=32,
        per_device_eval_batch_size=32,
        evaluation_strategy="steps",
        eval_steps=5_000,
        logging_steps=5_000,
        gradient_accumulation_steps=8,
        num_train_epochs=1,
        weight_decay=0.1,
        warmup_steps=1_000,
        lr_scheduler_type="cosine",
        learning_rate=5e-4,
        save_steps=5_000,
        fp16=True,
        push_to_hub=True,
    )

    trainer = Trainer(
        model=model,
        tokenizer=tokenizer,
        args=args,
        data_collator=data_collator,
        train_dataset=tokenized_datasets["train"],
        eval_dataset=tokenized_datasets["valid"],
    )

现在我们只需启动 `Trainer` 并等待训练完成。根据你是在整个训练集还是在训练集的一个子集上运行它，这将分别需要 20 或 2 个小时，因此请喝杯咖啡或者找一本好书来阅读！

    trainer.train()

训练完成后，我们可以将模型和 tokenizer 推送到 Hub：

    trainer.push_to_hub()

✏️ 试试看！ 除了 `TrainingArguments` 之外，我们只需要大约 30 行代码就可以从原始文本到训练 GPT-2。用你自己的数据集试试看，看看你能不能得到好的结果！

💡 如果你能使用多 GPU 的机器，尝试在那里运行代码。 `Trainer` 自动管理多台机器，这能极大地加快训练速度。

##
使用 pipeline 进行代码生成

现在是见证奇迹的时刻：我们来看看训练好的模型到底表现如何！我们可以在日志中看到损失持续下降，但要测试模型的效果，我们就看看它对一些提示的反应如何。为此，我们将模型包装在一个文本生成的 `pipeline` 中，并如果有 GPU 可用，我们将把它放在 GPU 上加快生成速度：

    import torch
    from transformers import pipeline

    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
    pipe = pipeline(
        "text-generation", model="huggingface-course/codeparrot-ds", device=device
    )

让我们从简单的创建散点图任务开始：

    txt = """\
    # 创建一些数据
    x = np.random.randn(100)
    y = np.random.randn(100)

    # 使用 x,y 创建散点图
    """
    print(pipe(txt, num_return_sequences=1)[0]["generated_text"])

    # 创建一些数据
    x = np.random.randn(100)
    y = np.random.randn(100)

    # 使用 x,y 创建散点图
    plt.scatter(x, y)

    # 创建散点

结果看起来是正确的。那么对于 `pandas` 操作也可以吗？让我们看看是否能从两个数组创建一个 `DataFrame` ：

    txt = """\
    # 创建一些数据
    x = np.random.randn(100)
    y = np.random.randn(100)

    # 从 x 和 y 创建 dataframe
    """
    print(pipe(txt, num_return_sequences=1)[0]["generated_text"])

    # 创建一些数据
    x = np.random.randn(100)
    y = np.random.randn(100)

    # 从 x 和 y 创建 dataframe
    df = pd.DataFrame({'x': x, 'y': y})
    df.insert(0,'x', x)
    for

很好，这是正确的答案——尽管它又把 `x` 重复插入了一次。而且由于生成的 token 数量有限，所以下面的 `for` 循环被切断了。让我们看看我们是否能做些更复杂的事情，让模型帮助我们使用 `groupby` 操作：

    txt = """\
    # 有职业,收入和名字的 dataframe
    df = pd.DataFrame({'profession': x, 'income':y, 'name': z})

    # 计算每个职业的平均收入
    """
    print(pipe(txt, num_return_sequences=1)[0]["generated_text"])

    # 有职业,收入和名字的 dataframe
    df = pd.DataFrame({'profession': x, 'income':y, 'name': z})

    # 计算每个职业的平均收入
    profession = df.groupby(['profession']).mean()

    # 计算

不错；是正确的。最后，让我们看看是否能引导模型使用 `scikit-learn` 并建立一个随机森林模型：

    txt = """
    # 从 scikit-learn 导入随机森林回归器
    from sklearn.ensemble import RandomForestRegressor

    # 用 X, y 拟合带有 300 个估算器的随机森林模型:
    """
    print(pipe(txt, num_return_sequences=1)[0]["generated_text"])

    # 从 scikit-learn 导入随机森林回归器
    from sklearn.ensemble import RandomForestRegressor

    # 用 X, y 拟合带有 300 个估算器的随机森林模型:
    rf = RandomForestRegressor(n_estimators=300, random_state=random_state, max_depth=3)
    rf.fit(X, y)
    rf

从这几个例子来看，模型似乎已经学习了 Python 数据科学堆栈的一些语法（当然，在将模型部署到现实世界之前，我们需要对其进行更全面的评估）。然而，有时候它需要更多的模型训练定制来达到特定情境所需的性能。例如，如果我们想动态更新 `batch_size` 或添加一个条件训练循环来跳过坏示例怎么办？一种选择是修改 `Trainer` 添加新的功能，但有时从头开始编写训练循环会更简单。这就是🤗 Accelerate 的用武之地。

##
使用🤗 Accelerate 进行训练

我们已经看到了如何使用 `Trainer` 训练模型，在 `Trainer` 中可以对训练过程可以通过修改一些参数进行一些定制。然而，有时我们想要完全控制训练循环，或者我们想要进行一些更自由的的更改。在这种情况下 🤗 Accelerate 是一个不错的选择，本节我们将介绍如何使用它来训练我们的模型。为了让事情变得更有趣，相比于上面的 `Trainer` 我们还将在训练循环中添加一些修改。

由于我们主要关注的是为数据科学库提供合理的代码自动补充功能，因此对于更多使用这些库的训练样本赋予更高的权重是有意义的。我们可以通过使用 `plt` 、 `pd` 、 `sk` 、 `fit` 和 `predict` 等关键词来轻松地识别出这些例子，这些关键词是 `matplotlib.pyplot` 、 `pandas` 和 `sklearn` 导入后最常用重命名的名称，以及 `sklearn` 的 `fit/predict` 方法。如果这些在模型的内部是用单一的一个 `token` 表示的，我们可以通过 token 的 id 轻松地检查它们是否出现在输入序列中。然而，Tokens 有可能有空格前缀，所以我们也需要在 tokenizer 词汇表中检查这些关键词。为了验证这个策略的有效性，我们会在测试样本中添加一个应该被分割为多个 tokens 的测试 token：

    keytoken_ids = []
    for keyword in [
        "plt",
        "pd",
        "sk",
        "fit",
        "predict",
        " plt",
        " pd",
        " sk",
        " fit",
        " predict",
        "testtest",
    ]:
        ids = tokenizer([keyword]).input_ids[0]
        if len(ids) == 1:
            keytoken_ids.append(ids[0])
        else:
            print(f"Keyword has not single token: {keyword}")

    'Keyword has not single token: testtest'

太好了，这个方法似乎很有效！我们现在可以编写一个自定义的损失函数，它的输入有输入序列、logits 和我们刚刚选择的关键字。首先需要对齐 `logits` 和 `inputs` : 并将输入序列右移一个单位形成目标序列，因为下一个 `token` 就是当前 `token` 的预测的目标。我们可以通过从输入序列的第二个 `token` 开始设置标签，因为模型不会预测第一个 `token`。然后我们截断最后一个 `logit`，因为我们没有完整输入序列后面的标签。有了这些，我们就可以计算每个样本的损失，并计算每个样本中所有关键词的出现次数。最后，我们使用出现次数作为权重，计算所有样本的加权平均值。由于我们不想抛弃所有没有关键词的样本，我们将所有的权重都加 1：

    from torch.nn import CrossEntropyLoss
    import torch

    def keytoken_weighted_loss(inputs, logits, keytoken_ids, alpha=1.0):
        # 左移 tokens < n 预测 n
        shift_labels = inputs[..., 1:].contiguous()
        shift_logits = logits[..., :-1, :].contiguous()
        # 计算每一个token的loss
        loss_fct = CrossEntropyLoss(reduce=False)
        loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
        # 对于每个样本重新调整大小并平均
        loss_per_sample = loss.view(shift_logits.size(0), shift_logits.size(1)).mean(axis=1)
        # 计算并缩放权重
        weights = torch.stack([(inputs == kt).float() for kt in keytoken_ids]).sum(
            axis=[0, 2]
        )
        weights = alpha * (1.0 + weights)
        # 计算评价权重
        weighted_loss = (loss_per_sample * weights).mean()
        return weighted_loss

在我们开始使用这个精妙的新损失函数进行训练之前，我们需要准备一些事情：

- 我们需要数据加载器来批量加载数据。
- 我们需要设置权重衰减参数。
- 有时我们在调试模型的时候可能需要临时评估，所以将评估代码包装在一个函数中。

让我们从数据加载器开始。我们只需要将数据集的格式设置为 `"torch"` ，然后我们就可以将它传递给一个具有适当 `batch size` 的 PyTorch 的 `DataLoader` ：

    from torch.utils.data.dataloader import DataLoader

    tokenized_datasets.set_format("torch")
    train_dataloader = DataLoader(tokenized_datasets["train"], batch_size=32, shuffle=True)
    eval_dataloader = DataLoader(tokenized_datasets["valid"], batch_size=32)

接下来，我们将参数分组，以便优化器知道哪些参数需要进行额外的权重衰减。通常，所有的偏置和 LayerNorm 权重项都不需要进行权重衰减；因此我们可以这样做：

    weight_decay = 0.1

    def get_grouped_params(model, no_decay=["bias", "LayerNorm.weight"]):
        params_with_wd, params_without_wd = [], []
        for n, p in model.named_parameters():
            if any(nd in n for nd in no_decay):
                params_without_wd.append(p)
            else:
                params_with_wd.append(p)
        return [
            {"params": params_with_wd, "weight_decay": weight_decay},
            {"params": params_without_wd, "weight_decay": 0.0},
        ]

我们希望在训练过程中定期在验证集上评估模型，让我们为此编写一个函数。它只需遍历评估数据加载器，并收集所有进程中的损失值：

    def evaluate():
        model.eval()
        losses = []
        for step, batch in enumerate(eval_dataloader):
            with torch.no_grad():
                outputs = model(batch["input_ids"], labels=batch["input_ids"])

            losses.append(accelerator.gather(outputs.loss))
        loss = torch.mean(torch.cat(losses))
        try:
            perplexity = torch.exp(loss)
        except OverflowError:
            perplexity = float("inf")
        return loss.item(), perplexity.item()

通过 `evaluate()` 函数我们定期可以获取损失值和 困惑度（perplexity） 。接下来，我们重新加载我们的模型以确保我们再次从头开始训练，而不是从上面的 `Trainer` 继续微调：

    model = GPT2LMHeadModel(config)

然后我们可以定义我们的优化器，使用之前的函数来分割权重衰减的参数：

    from torch.optim import AdamW

    optimizer = AdamW(get_grouped_params(model), lr=5e-4)

现在让我们准备模型、优化器和数据加载器，然后我们可以开始训练：

    from accelerate import Accelerator

    accelerator = Accelerator(fp16=True)

    model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(
        model, optimizer, train_dataloader, eval_dataloader
    )

🚨 如果你在 TPU 上训练，你需要将上述单元格开始的所有代码移到一个专门的训练函数中。更多详情请参阅 第三章 。

现在我们已经将我们的 `train_dataloader` 传递给了 `accelerator.prepare()` ，我们可以使用 `len()` 来计算训练步骤的数量。请记住，我们应该在准备好 `dataloader` 后再使用 `len()` ，因为改动 `dataloader` 会改变其长度。我们使用一个从学习率衰减到 0 的经典线性学习率调度：

    num_train_epochs = 1
    num_update_steps_per_epoch = len(train_dataloader)
    num_training_steps = num_train_epochs * num_update_steps_per_epoch

    lr_scheduler = get_scheduler(
        name="linear",
        optimizer=optimizer,
        num_warmup_steps=1_000,
        num_training_steps=num_training_steps,
    )

最后，为了将我们的模型推送到 Hub，我们需要在一个工作文件夹中创建一个 `Repository` 对象。如果你还没有登录的话，首先需要登录到 Hugging Face，我们将根据模型 ID 来确定仓库名称（你可以使用你喜欢的名字替换 `repo_name` ；它只需要包含你的用户名，可以使用 `get_full_repo_name()` 函数的查看目前的 `repo_name`）：

    from huggingface_hub import Repository, get_full_repo_name

    model_name = "codeparrot-ds-accelerate"
    repo_name = get_full_repo_name(model_name)
    repo_name

    'sgugger/codeparrot-ds-accelerate'

然后我们可以将该仓库克隆到本地文件夹中。如果本地已经存在一个同名的文件夹，这个本地文件夹应该是我们正在使用的仓库的克隆在本地的版本：

    output_dir = "codeparrot-ds-accelerate"
    repo = Repository(output_dir, clone_from=repo_name)

我们现在可以通过调用 `repo.push_to_hub()` 方法上传保存在 `output_dir` 中的所有内容。这将帮助我们在每个训练周期结束时上传中间模型。

    evaluate()

    (10.934126853942871, 56057.14453125)

目前的损失和困惑度都是非常高的值，但这并不奇怪，因为我们还没有训练模型。到现在为止，我们已经为编写训练脚本的核心部分：训练循环已经做好了准备。在训练循环中，我们迭代遍历数据加载器并将成批量的数据传递给模型。有了模型输出的 logits，我们就可以使用自定义损失函数计算损伤。我们通过梯度累积步骤的数量来缩放损失，以避免在聚合更多步骤时产生更大的损失。在我们优化之前，我们也会剪裁梯度来更好的收敛。最后，每隔一段步数，我们用新的 `evaluate()` 函数在评估集上评估模型：

    from tqdm.notebook import tqdm

    gradient_accumulation_steps = 8
    eval_steps = 5_000

    model.train()
    completed_steps = 0
    for epoch in range(num_train_epochs):
        for step, batch in tqdm(
            enumerate(train_dataloader, start=1), total=num_training_steps
        ):
            logits = model(batch["input_ids"]).logits
            loss = keytoken_weighted_loss(batch["input_ids"], logits, keytoken_ids)
            if step % 100 == 0:
                accelerator.print(
                    {
                        "samples": step * samples_per_step,
                        "steps": completed_steps,
                        "loss/train": loss.item() * gradient_accumulation_steps,
                    }
                )
            loss = loss / gradient_accumulation_steps
            accelerator.backward(loss)
            if step % gradient_accumulation_steps == 0:
                accelerator.clip_grad_norm_(model.parameters(), 1.0)
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()
                completed_steps += 1
            if (step % (eval_steps * gradient_accumulation_steps)) == 0:
                eval_loss, perplexity = evaluate()
                accelerator.print({"loss/eval": eval_loss, "perplexity": perplexity})
                model.train()
                accelerator.wait_for_everyone()
                unwrapped_model = accelerator.unwrap_model(model)
                unwrapped_model.save_pretrained(output_dir, save_function=accelerator.save)
                if accelerator.is_main_process:
                    tokenizer.save_pretrained(output_dir)
                    repo.push_to_hub(
                        commit_message=f"Training in progress step {step}", blocking=False
                    )

就是这样 - 你现在拥有自己的因果语言模型（例如 GPT-2）的自定义训练循环，你可以根据自己的需要进一步定制。

✏️ 试试看！ 创建适合你的用例的自定义损失函数，或在训练循环中添加另一个自定义步骤。

✏️ 试试看！ 当运行长时间的训练实验时，使用 TensorBoard 或 Weights & Biases 等工具记录重要指标是个好主意。向训练循环中添加适当的日志记录，这样你可以随时检查训练进度。

< > Update on GitHub