从零开始复现GPT2（五）：评估和可视化代码的实现

---

源码地址：https://gitee.com/guojialiang2023/gpt2

---

模型

评估代码

配置

from typing import Iterator

class EvaluateConfig(object):
    def __init__(self,
                 batch_eval: int,
                 total_steps: int,
                 use_gpu: bool):
        self.batch_eval = batch_eval
        self.total_steps = total_steps
        self.use_gpu = use_gpu

    def iterate(self) -> Iterator:
        if self.total_steps == -1:
            while True:
                yield
        else:
            yield from range(self.total_steps)

评估框架

import torch
import torch.nn as nn
from data import Dataset
from typing import Dict


class EvaluationSpec(object):
    def initialize(self):
        pass

    def prepare_dataset(self) -> Dataset:
        raise NotImplementedError()

    def construct_model(self) -> nn.Module:
        raise NotImplementedError()

    def eval_objective(self, data: Dict[str, torch.Tensor], model: nn.Module
                       ) -> Dict[str, torch.Tensor]:
        raise NotImplementedError()

评估类实现

代码定义了一个用于模型评估的类Evaluator，主要用于在给定的评估规范(EvaluationSpec)和配置(EvaluateConfig)下，对模型进行性能评估。以下是详细的解释：

Evaluator类定义

• __init__(self, spec: EvaluationSpec, config: EvaluateConfig)：构造函数，接收两个参数spec和config。spec是一个EvaluationSpec对象，定义了评估过程中的各种操作，如如何初始化环境、准备数据集、构建模型以及计算评估指标等。config是一个EvaluateConfig对象，包含评估过程的配置信息，如是否使用GPU、批量大小等。

evaluate方法

• evaluate(self, from_model: Optional[str] = None) -> Dict[str, float]：评估方法，可选地接收一个模型文件路径from_model，用于加载预训练的模型参数。该方法执行评估过程，并返回一个字典，键为评估指标的名称，值为对应的平均分数。

1. 初始化和数据集准备：通过调用spec.initialize()初始化评估环境，并通过spec.prepare_dataset()准备评估所需的数据集。

2. 模型加载：构建模型实例并设置为评估模式。如果提供了模型路径(from_model)，则从该路径加载模型权重。

3. 设备和精度设置：如果配置中指定使用GPU(config.use_gpu)，则将模型转移到GPU上，并将模型的数据类型转换为半精度浮点数，以加快计算速度并减少内存消耗。

4. 评估循环：使用config.iterate()方法迭代地进行评估。对于每一批数据，通过_eval_step方法计算评估指标。收集每批的评估指标，并计算所有批次的平均值。

_eval_step方法

• _eval_step(self, dataset: Dataset, model: nn.Module) -> Optional[Dict[str, float]]：私有方法，执行单批次的评估步骤。接收数据集和模型作为参数，计算并返回当前批次的评估指标。如果数据集已遍历完毕，返回None。

1. 数据获取和处理：从数据集中获取一批数据，如果配置中指定使用GPU，将数据转移到GPU上。

2. 评估指标计算：调用spec.eval_objective方法，传入当前批次的数据和模型，计算评估指标。

3. 异常处理：使用try-except捕获StopIteration异常，这表示数据集已经遍历完毕，此时方法返回None。

import torch
import torch.nn as nn
from data import Dataset
from evaluation import EvaluationSpec, EvaluateConfig
from typing import Optional, Dict


class Evaluator(object):
    def __init__(self, spec: EvaluationSpec, config: EvaluateConfig):
        self.spec = spec
        self.config = config

    def evaluate(self, from_model: Optional[str] = None) -> Dict[str, float]:
        # Initialize evaluation environment and prepare a dataset.
        self.spec.initialize()
        eval_dataset = self.spec.prepare_dataset()

        # Load trained model parameters.
        model = self.spec.construct_model().eval()
        if from_model:
            ckpt = torch.load(from_model, map_location='cpu')
            model.load_state_dict(ckpt['model'])

        # Move the model to GPU device and convert the data type to half
        # precision.
        if self.config.use_gpu:
            model.cuda().half()

        total_metrics = {}
        for _ in self.config.iterate():
            batch_metrics = self._eval_step(eval_dataset, model)
            if batch_metrics is None:
                break

            # Record the batched metrics.
            for k, v in batch_metrics.items():
                if k not in total_metrics:
                    total_metrics[k] = []
                total_metrics[k].append(v)

        return {k: sum(v) / len(v) for k, v in total_metrics.items()}

    @torch.no_grad()
    def _eval_step(self, dataset: Dataset, model: nn.Module
                   ) -> Optional[Dict[str, float]]:
        try:
            data = dataset.fetch(self.config.batch_eval)
            if self.config.use_gpu:
                data = {k: v.cuda() for k, v in data.items()}

            metrics = self.spec.eval_objective(data, model)
            return {k: v.item() for k, v in metrics.items()}
        except StopIteration:
            return None

评估代码

代码展示了如何为GPT-2模型的评估过程定义一个特定的评估规范（GPT2EvaluationSpec），并通过argparse库接收命令行参数来执行模型评估。整个过程分为几个关键部分：

GPT2EvaluationSpec类定义

• 构造函数：初始化评估规范时，接收评估数据集、词汇表路径、序列长度、模型的层数、头数、维度以及维度增长率等参数。

• initialize方法：加载词汇表并初始化损失函数（交叉熵损失）。

• prepare_dataset方法：使用指定的评估语料库和词汇表创建一个TokenizedCorpus实例，这是一个经过预处理的、为评估准备好的数据集。

• construct_model方法：根据提供的模型配置参数构造一个Transformer模型实例。

• eval_objective方法：计算评估指标，包括损失和困惑度（perplexity）。这通过将模型的输出和预期输出进行对比来实现，同时忽略填充（pad）标记的损失贡献。

evaluate_gpt2_model函数

• 参数解析：接收通过命令行传入的参数，包括模型路径、评估语料库路径、词汇表路径、模型配置（序列长度、层数、头数、维度等）以及评估配置（批量大小、总步数、是否使用GPU）。

• 评估执行：创建GPT2EvaluationSpec和EvaluateConfig实例，然后实例化Evaluator并调用其evaluate方法来执行评估，最后打印评估结果。

add_subparser函数

• 子命令配置：定义了一个子命令evaluate，用于通过命令行界面配置和执行GPT-2模型的评估。这包括模型路径、评估相关路径、模型配置以及评估选项的具体参数。

整体流程

1. 配置解析：通过命令行接口接收用户输入的参数。
2. 评估准备：根据接收到的参数，准备评估所需的资源，包括数据集、模型、词汇表等。
3. 模型评估：执行模型评估，计算损失和困惑度等评估指标。
4. 结果输出：将评估结果输出到控制台。

这个实现体现了在NLP模型开发中常见的模式，即通过特定的评估规范和灵活的命令行参数接收方式，来支持模型的评估和性能分析。通过这种方式，研究者和开发者可以轻松地调整模型参数和评估设置，以适应不同的需求和实验场景。

import argparse
import torch
import torch.nn as nn
from modeling import Transformer
from data import Dataset, Vocab, TokenizedCorpus
from evaluation import EvaluationSpec, EvaluateConfig, Evaluator
from typing import Dict


class GPT2EvaluationSpec(EvaluationSpec):
    def __init__(self, eval_corpus: str, vocab_path: str, seq_len: int,
                 layers: int, heads: int, dims: int, rate: int):
        self.eval_corpus = eval_corpus
        self.vocab_path = vocab_path
        self.seq_len = seq_len
        self.layers = layers
        self.heads = heads
        self.dims = dims
        self.rate = rate

    def initialize(self):
        self.vocab = Vocab(vocab_path=self.vocab_path)
        self.criterion = nn.CrossEntropyLoss(reduction='none')

    def prepare_dataset(self) -> Dataset:
        return TokenizedCorpus(corpus_path=self.eval_corpus,
                               vocab=self.vocab,
                               seq_len=self.seq_len,
                               repeat=False)

    def construct_model(self) -> nn.Module:
        return Transformer(layers=self.layers, pad_idx=self.vocab.pad_idx,
                           words=len(self.vocab), seq_len=self.seq_len,
                           heads=self.heads, dims=self.dims, rate=self.rate,
                           dropout=0, bidirectional=False)

    def eval_objective(self, data: Dict[str, torch.Tensor], model: nn.Module
                       ) -> Dict[str, torch.Tensor]:
        logits, _ = model(data['input'], past=None)
        loss = self.criterion(logits.transpose(1, 2), data['output'])

        mask = (data['output'] != self.vocab.pad_idx).float()
        loss = (loss * mask).sum() / mask.sum()
        perplexity = (loss.exp() * mask).sum() / mask.sum()

        return {'loss': loss, 'perplexity': perplexity}


def evaluate_gpt2_model(args: argparse.Namespace):
    spec = GPT2EvaluationSpec(
        eval_corpus=args.eval_corpus, vocab_path=args.vocab_path,
        seq_len=args.seq_len, layers=args.layers, heads=args.heads,
        dims=args.dims, rate=args.rate)
    config = EvaluateConfig(
        batch_eval=args.batch_eval, total_steps=args.total_steps,
        use_gpu=args.use_gpu)

    print(Evaluator(spec, config).evaluate(from_model=args.model_path))


def add_subparser(subparsers: argparse._SubParsersAction):
    parser = subparsers.add_parser('evaluate', help='evaluate GPT-2 model')

    parser.add_argument('--model_path', required=True,
                        help='trained GPT-2 model file path')

    group = parser.add_argument_group('Corpus and vocabulary')
    group.add_argument('--eval_corpus', required=True,
                       help='evaluation corpus file path')
    group.add_argument('--vocab_path', required=True,
                       help='vocabulary file path')

    group = parser.add_argument_group('Model configurations')
    group.add_argument('--seq_len', default=64, type=int,
                       help='maximum sequence length')
    group.add_argument('--layers', default=12, type=int,
                       help='number of transformer layers')
    group.add_argument('--heads', default=16, type=int,
                       help='number of multi-heads in attention layer')
    group.add_argument('--dims', default=1024, type=int,
                       help='dimension of representation in each layer')
    group.add_argument('--rate', default=4, type=int,
                       help='increase rate of dimensionality in bottleneck')

    group = parser.add_argument_group('Evaluation options')
    group.add_argument('--batch_eval', default=64, type=int,
                       help='number of evaluation batch size')
    group.add_argument('--total_steps', default=-1, type=int,
                       help='number of total evaluation steps')
    group.add_argument('--use_gpu', action='store_true',
                       help='use gpu device in inferencing')

    parser.set_defaults(func=evaluate_gpt2_model)

可视化代码

代码定义了一个用于可视化训练和评估过程中记录的指标的工具，特别是针对损失函数的值。它包括四种不同的图表来展示这些指标，每种图表侧重于不同方面的展示。此工具使用matplotlib库来生成图表，并通过argparse库接受命令行参数来控制其行为。下面是详细的解释：

_plot_entire_metrics_graph 函数

• 这个函数绘制了训练和评估损失随迭代次数变化的整体趋势图。它使用不同的颜色来区分训练损失和评估损失，提供了一个全局视角来观察模型的性能。

_plot_log_scale_metrics_graph 函数

• 这个函数也绘制了训练和评估损失，但是将x轴（迭代次数）设置为对数尺度。这种方式特别适用于指数级变化的数据，可以帮助观察在训练初期损失迅速下降的行为。

_plot_stretched_metrics_graph 函数

• 这个函数专注于展示训练和评估损失在最近90%的迭代中的详细变化。它通过调整y轴的范围来放大这部分数据，从而更清晰地观察模型在训练后期的性能表现。

_plot_highlight_metrics_graph 函数

• 这个函数采用了平滑处理来展示评估损失，并只展示最后30%的迭代数据。它旨在通过减少评估损失的波动，更清楚地展示趋势。同时，它也在图表中保留了原始的评估损失数据，以便进行比较。

visualize_recorded_metrics 函数

• 这个函数是主要的入口点，用于加载模型训练过程中记录的指标，并调用上述四个函数将这些指标可视化。它支持两种模式：交互式和非交互式。在交互式模式下，会直接显示图表窗口；在非交互式模式下，会将图表保存为文件。

add_subparser 函数

• 这个函数用于在命令行工具中添加一个子命令visualize，允许用户指定模型文件路径、输出图像文件的路径以及是否启用交互式模式。

整体上，这个工具提供了一种灵活的方式来分析和理解模型在训练和评估过程中的性能变化，对于模型调优和理解其行为模式非常有帮助。通过不同的图表展示方法，用户可以从多个角度审视损失函数的变化，帮助识别过拟合、欠拟合以及其他可能的问题。

import math
import torch
import argparse
import matplotlib.pyplot as plt
from typing import List


def _plot_entire_metrics_graph(train_steps: List[int],
                               train_losses: List[float],
                               eval_steps: List[int],
                               eval_losses: List[float]):
    plt.plot(train_steps, train_losses,
             label='train', color='#4dacfa', linewidth=2, zorder=10)
    plt.plot(eval_steps, eval_losses,
             label='evaluate', color='#ff6e54', linewidth=2, zorder=0)

    plt.title('Cross-Entropy Loss')
    plt.xlabel('Iterations')
    plt.ylabel('Loss')

    plt.legend(loc='upper right')
    plt.gca().xaxis.set_major_formatter(
        plt.FuncFormatter(lambda x, pos: f'{math.floor(x / 1000)}k'))


def _plot_log_scale_metrics_graph(train_steps: List[int],
                                  train_losses: List[float],
                                  eval_steps: List[int],
                                  eval_losses: List[float]):
    plt.plot(train_steps, train_losses,
             label='train', color='#4dacfa', linewidth=2, zorder=10)
    plt.plot(eval_steps, eval_losses,
             label='evaluate', color='#ff6e54', linewidth=2, zorder=0)

    plt.title('Log-Scale Loss Graph')
    plt.xlabel('Iterations (Log Scale)')
    plt.ylabel('Loss')

    plt.xscale('log')

    plt.legend(loc='upper right')
    plt.gca().xaxis.set_major_formatter(
        plt.FuncFormatter(lambda x, pos: f'{math.floor(x / 1000)}k'))


def _plot_stretched_metrics_graph(train_steps: List[int],
                                  train_losses: List[float],
                                  eval_steps: List[int],
                                  eval_losses: List[float]):
    target_range_train = len(train_steps) * 9 // 10
    target_range_eval = len(eval_steps) * 9 // 10
    min_loss = min(train_losses[-target_range_train:]
                   + eval_losses[-target_range_eval:])
    max_loss = max(train_losses[-target_range_train:]
                   + eval_losses[-target_range_eval:])

    plt.plot(train_steps, train_losses,
             label='train', color='#4dacfa', linewidth=2, zorder=10)
    plt.plot(eval_steps, eval_losses,
             label='evaluate', color='#ff6e54', linewidth=1, zorder=1)

    plt.title('Detail Loss Graph')
    plt.xlabel('Iterations')
    plt.ylabel('Loss')

    plt.ylim(min_loss - 0.1, max_loss + 0.1)

    plt.legend(loc='upper right')
    plt.gca().xaxis.set_major_formatter(
        plt.FuncFormatter(lambda x, pos: f'{math.floor(x / 1000)}k'))


def _plot_highlight_metrics_graph(train_steps: List[int],
                                  train_losses: List[float],
                                  eval_steps: List[int],
                                  eval_losses: List[float]):
    beta = 2 / (100 + 1)
    smoothed_eval_losses = [eval_losses[0]]
    for loss in eval_losses[1:]:
        smoothed_eval_losses.append(
            beta * loss + (1 - beta) * smoothed_eval_losses[-1])

    target_range_train = len(train_steps) * 3 // 10
    target_range_eval = len(eval_steps) * 3 // 10

    plt.plot(train_steps[-target_range_train:],
             train_losses[-target_range_train:],
             label='train', color='#4dacfa', linewidth=2, zorder=10)
    plt.plot(eval_steps[-target_range_eval:],
             smoothed_eval_losses[-target_range_eval:],
             label='evaluate', color='#ff6e54', linewidth=2, zorder=1)
    plt.plot(eval_steps[-target_range_eval:],
             eval_losses[-target_range_eval:],
             color='#ff6e543f', linewidth=3, zorder=0)

    plt.title('Detail Loss Graph')
    plt.xlabel('Iterations')
    plt.ylabel('Loss')

    plt.legend(loc='upper right')
    plt.gca().xaxis.set_major_formatter(
        plt.FuncFormatter(lambda x, pos: f'{math.floor(x / 1000)}k'))


def visualize_recorded_metrics(args: argparse.Namespace):
    metrics = torch.load(args.model_path)['metrics']
    train_steps, train_losses = zip(*metrics['train/loss'])
    eval_steps, eval_losses = zip(*metrics['eval/loss'])

    plt.figure(figsize=(10, 7))

    plt.subplot(221)
    _plot_entire_metrics_graph(train_steps, train_losses,
                               eval_steps, eval_losses)

    plt.subplot(222)
    _plot_log_scale_metrics_graph(train_steps, train_losses,
                                  eval_steps, eval_losses)

    plt.subplot(223)
    _plot_stretched_metrics_graph(train_steps, train_losses,
                                  eval_steps, eval_losses)

    plt.subplot(224)
    _plot_highlight_metrics_graph(train_steps, train_losses,
                                  eval_steps, eval_losses)

    plt.tight_layout()
    if args.interactive:
        plt.show()
    else:
        plt.savefig(args.figure)


def add_subparser(subparsers: argparse._SubParsersAction):
    parser = subparsers.add_parser(
        'visualize', help='visualize metrics recorded during training')

    parser.add_argument('--figure', default='figure.png',
                        help='output figure image file path')
    parser.add_argument('--model_path', required=True,
                        help='trained GPT-2 model file path')
    parser.add_argument('--interactive', action='store_true',
                        help='show interactive plot window')

    parser.set_defaults(func=visualize_recorded_metrics)