bertviz - 利用BERT、GPT2等模型进行注意力机制可视化

BertViz 项目介绍

项目概述

BertViz 是一个交互式工具，旨在可视化自然语言处理（NLP）模型中的注意力机制。它支持在 Jupyter 或 Colab 笔记本中运行，并提供了简单的Python API，兼容大多数 Huggingface 的模型，如 BERT、GPT-2、T5 等。BertViz 扩展了 Tensor2Tensor 的可视化工具，提供了多种视图来深入探索注意力机制。

快速体验

BertViz 提供了三种主要的视图方式，让用户能从不同的角度观察和理解模型中的注意力机制。

Head View

Head View 可视化同一层中一个或多个注意头的注意力。用户可以在示例代码中选择不同的模型和输入文本以查看这些可视化。

Model View

Model View 提供了从高空俯瞰的视角，展示了跨多个层和头的注意力分布。

Neuron View

Neuron View 让用户可视化个别神经元的查询和键向量，并展示其如何被用来计算注意力。

如何开始

在 Jupyter Notebook 中运行 BertViz

用户可以通过以下命令行命令安装 BertViz 以及所需的 Jupyter Notebook 和 ipywidgets：

pip install bertviz
pip install jupyterlab
pip install ipywidgets

安装完毕后，运行 jupyter notebook 创建一个新的 Jupyter notebook，并选择 Python 3 内核。

在 Colab 上运行 BertViz

在 Colab 上，只需在笔记本开始处添加以下命令即可：

!pip install bertviz

样例代码

以下是一个加载模型并在 Model View 中展示注意力的样例代码：

from transformers import AutoTokenizer, AutoModel, utils
from bertviz import model_view
utils.logging.set_verbosity_error()

model_name = "microsoft/xtremedistil-l12-h384-uncased"
input_text = "The cat sat on the mat"
model = AutoModel.from_pretrained(model_name, output_attentions=True)
tokenizer = AutoTokenizer.from_pretrained(model_name)
inputs = tokenizer.encode(input_text, return_tensors='pt')
outputs = model(inputs)
attention = outputs[-1]
tokens = tokenizer.convert_ids_to_tokens(inputs[0])
model_view(attention, tokens)

用户可以更换不同的输入文本和模型来尝试不同的可视化效果。

互动教程

BertViz 提供了一个互动教程，用户可以体验和学习如何利用这个工具进行注意力的可视化。

文档与支持

BertViz 提供详细的文档，包含每种视图的使用说明、自注意模型与编码解码模型的支持，以及各种附加选项和功能限制等信息。

结论

BertViz 是一个功能强大的工具，帮助研究人员和工程师深入理解和分析 Transformer 模型中的注意力机制。通过多种视图和简单的代码接口，用户可以对模型的内部工作原理有更直观的理解。