Logo
Project Icon

nnsight

解释和操作深度学习模型内部的Python包

nnsight深度学习模型神经网络解释模型操作PyTorchGithub开源项目

nnsight是一个专门用于深度学习模型内部解释和操作的Python包。它可以访问模型隐藏状态、进行噪声注入和跨提示干预。该工具支持保存中间值、修改参数和多token生成等功能，方便研究人员和开发者深入分析和调试神经网络模型。

访问官网访问官网GithubGithub论文论文
介绍相关项目
drawing

nnsight

nnsight 包允许解释和操作深度学习模型的内部结构。阅读我们的论文！

安装

通过pip安装此包，运行：

pip install nnsight

示例

这是一个简单示例，我们在本地对gpt2运行nnsight API并保存最后一层的隐藏状态：

from nnsight import LanguageModel

model = LanguageModel('openai-community/gpt2', device_map='auto')

with model.trace('The Eiffel Tower is in the city of'):

      hidden_states = model.transformer.h[-1].output[0].save()

      output = model.output.save()

让我们逐步分析。

我们从nnsight模块导入LanguageModel对象，并使用gpt2的huggingface仓库ID 'openai-community/gpt2'创建一个gpt2模型。这接受创建模型的参数，包括device_map来指定运行的设备。

from nnsight import LanguageModel

model = LanguageModel('openai-community/gpt2',device_map='auto')

然后，我们通过在模型对象上调用.trace(...)来创建一个追踪上下文块。这表示我们要用我们的提示运行模型。

with model.trace('The Eiffel Tower is in the city of') as tracer:

现在调用.trace(...)实际上并不初始化或运行模型。只有在退出追踪块后，才会加载和运行实际模型。块中的所有操作都是"代理"，本质上创建了我们希望稍后执行的操作图。

在这个上下文中，所有操作/干预将应用于给定提示的处理。

hidden_states = model.transformer.h[-1].output[0].save()

在这一行中，我们说，访问transformer的最后一层model.transformer.h[-1]，访问其输出.output，索引为0 .output[0]，并保存它.save()

有几点需要注意，我们可以通过打印模型来查看模型的模块树。这允许我们知道要访问哪些属性来获取所需的模块。 运行print(model)得到：

GPT2LMHeadModel(
  (transformer): GPT2Model(
    (wte): Embedding(50257, 768)
    (wpe): Embedding(1024, 768)
    (drop): Dropout(p=0.1, inplace=False)
    (h): ModuleList(
      (0-11): 12 x GPT2Block(
        (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
        (attn): GPT2Attention(
          (c_attn): Conv1D()
          (c_proj): Conv1D()
          (attn_dropout): Dropout(p=0.1, inplace=False)
          (resid_dropout): Dropout(p=0.1, inplace=False)
        )
        (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
        (mlp): GPT2MLP(
          (c_fc): Conv1D()
          (c_proj): Conv1D()
          (act): NewGELUActivation()
          (dropout): Dropout(p=0.1, inplace=False)
        )
      )
    )
    (ln_f): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
  )
  (lm_head): Linear(in_features=768, out_features=50257, bias=False)
)

.output返回此模块输出的代理。这本质上意味着我们在说，当我们在推理过程中到达这个模块的输出时，抓取它并执行我们在其上定义的任何操作（这些操作也成为代理）。这里有两个操作代理，一个用于获取输出的第0个索引，另一个用于保存输出。我们取第0个索引是因为gpt2 transformer层的输出是一个元组，其中第一个索引是实际的隐藏状态（最后两个索引来自注意力）。我们可以在任何代理上调用.shape来获取值最终的形状。 运行print(model.transformer.h[-1].output.shape)返回(torch.Size([1, 10, 768]), (torch.Size([1, 12, 10, 64]), torch.Size([1, 12, 10, 64])))

在处理我们正在构建的干预计算图时，当不再需要代理的值时，其值会被取消引用并销毁。然而，在代理上调用.save()会通知计算图保存此代理的值并永不销毁它，允许我们在生成后访问该值。

退出生成器上下文后，模型将使用指定的参数和干预图运行。output将填充实际输出，hidden_states将包含隐藏值。

print(output)
print(hidden_states)

返回：

tensor([[ 464,  412,  733,  417, 8765,  318,  287,  262, 1748,  286, 6342]],
       device='cuda:0')
tensor([[[ 0.0505, -0.1728, -0.1690,  ..., -1.0096,  0.1280, -1.0687],
         [ 8.7494,  2.9057,  5.3024,  ..., -8.0418,  1.2964, -2.8677],
         [ 0.2960,  4.6686, -3.6642,  ...,  0.2391, -2.6064,  3.2263],
         ...,
         [ 2.1537,  6.8917,  3.8651,  ...,  0.0588, -1.9866,  5.9188],
         [-0.4460,  7.4285, -9.3065,  ...,  2.0528, -2.7946,  0.5556],
         [ 6.6286,  1.7258,  4.7969,  ...,  7.6714,  3.0682,  2.0481]]],
       device='cuda:0')
操作

大多数基本操作和torch操作都适用于代理，并被添加到计算图中。

from nnsight import LanguageModel
import torch 

model = LanguageModel('openai-community/gpt2', device_map='cuda')

with model.trace('The Eiffel Tower is in the city of'):

  hidden_states_pre = model.transformer.h[-1].output[0].save()

  hs_sum = torch.sum(hidden_states_pre).save()

  hs_edited = hidden_states_pre + hs_sum

  hs_edited = hs_edited.save()

print(hidden_states_pre)
print(hs_sum)
print(hs_edited)

在这个例子中,我们获取隐藏状态的总和并将其添加到隐藏状态本身(无论出于什么原因)。通过保存各个步骤,我们可以看到值是如何变化的。

tensor([[[ 0.0505, -0.1728, -0.1690,  ..., -1.0096,  0.1280, -1.0687],
         [ 8.7494,  2.9057,  5.3024,  ..., -8.0418,  1.2964, -2.8677],
         [ 0.2960,  4.6686, -3.6642,  ...,  0.2391, -2.6064,  3.2263],
         ...,
         [ 2.1537,  6.8917,  3.8651,  ...,  0.0588, -1.9866,  5.9188],
         [-0.4460,  7.4285, -9.3065,  ...,  2.0528, -2.7946,  0.5556],
         [ 6.6286,  1.7258,  4.7969,  ...,  7.6714,  3.0682,  2.0481]]],
       device='cuda:0')
tensor(501.2957, device='cuda:0')
tensor([[[501.3461, 501.1229, 501.1267,  ..., 500.2860, 501.4237, 500.2270],
         [510.0451, 504.2014, 506.5981,  ..., 493.2538, 502.5920, 498.4279],
         [501.5916, 505.9643, 497.6315,  ..., 501.5348, 498.6892, 504.5219],
         ...,
         [503.4493, 508.1874, 505.1607,  ..., 501.3545, 499.3091, 507.2145],
         [500.8496, 508.7242, 491.9892,  ..., 503.3485, 498.5010, 501.8512],
         [507.9242, 503.0215, 506.0926,  ..., 508.9671, 504.3639, 503.3438]]],
       device='cuda:0')
设置

我们通常不仅想看到计算过程中发生的事情,还想干预和编辑信息流。

from nnsight import LanguageModel
import torch 

model = LanguageModel('openai-community/gpt2', device_map='cuda')

with model.trace('The Eiffel Tower is in the city of') as tracer:

  hidden_states_pre = model.transformer.h[-1].mlp.output.clone().save()

  noise = (0.001**0.5)*torch.randn(hidden_states_pre.shape)

  model.transformer.h[-1].mlp.output = hidden_states_pre + noise

  hidden_states_post = model.transformer.h[-1].mlp.output.save()

print(hidden_states_pre)
print(hidden_states_post)

在这个例子中,我们创建一个噪声张量来添加到隐藏状态中。然后我们添加它,使用赋值=运算符用这些新的带噪声的激活来更新.output的值。

我们可以在结果中看到变化:

tensor([[[ 0.0505, -0.1728, -0.1690,  ..., -1.0096,  0.1280, -1.0687],
         [ 8.7494,  2.9057,  5.3024,  ..., -8.0418,  1.2964, -2.8677],
         [ 0.2960,  4.6686, -3.6642,  ...,  0.2391, -2.6064,  3.2263],
         ...,
         [ 2.1537,  6.8917,  3.8651,  ...,  0.0588, -1.9866,  5.9188],
         [-0.4460,  7.4285, -9.3065,  ...,  2.0528, -2.7946,  0.5556],
         [ 6.6286,  1.7258,  4.7969,  ...,  7.6714,  3.0682,  2.0481]]],
       device='cuda:0')
tensor([[[ 0.0674, -0.1741, -0.1771,  ..., -0.9811,  0.1972, -1.0645],
         [ 8.7080,  2.9067,  5.2924,  ..., -8.0253,  1.2729, -2.8419],
         [ 0.2611,  4.6911, -3.6434,  ...,  0.2295, -2.6007,  3.2635],
         ...,
         [ 2.1859,  6.9242,  3.8666,  ...,  0.0556, -2.0282,  5.8863],
         [-0.4568,  7.4101, -9.3698,  ...,  2.0630, -2.7971,  0.5522],
         [ 6.6764,  1.7416,  4.8027,  ...,  7.6507,  3.0754,  2.0218]]],
       device='cuda:0')
多个令牌生成

当生成多个令牌时,使用.generate(...)和.next()来表示后续的干预应该应用于后续的生成。

这里我们再次使用gpt2生成,但生成三个令牌并保存最后一层的隐藏状态:

from nnsight import LanguageModel

model = LanguageModel('openai-community/gpt2', device_map='cuda')

with model.generate('埃菲尔铁塔位于', max_new_tokens=3) as tracer:
 
  hidden_states1 = model.transformer.h[-1].output[0].save()

  invoker.next()

  hidden_states2 = model.transformer.h[-1].next().output[0].save()

  invoker.next()

  hidden_states3 = model.transformer.h[-1].next().output[0].save()
跨提示干预

干预操作可以跨提示工作!在同一个生成块中使用两个调用,操作可以在它们之间工作。

你可以通过不向.trace/.generate传递提示来做到这一点,而是在创建的tracer对象上调用.invoke(...)。

在这种情况下,我们获取第一个提示"麦迪逊广场花园位于纽约"的令牌嵌入,并用它们替换第二个提示的嵌入。

from nnsight import LanguageModel

model = LanguageModel('openai-community/gpt2', device_map='cuda')

with model.generate(max_new_tokens=3) as tracer:
    
    with tracer.invoke("麦迪逊广场花园位于纽约"):

        embeddings = model.transformer.wte.output

    with tracer.invoke("_ _ _ _ _ _ _ _ _ _"):

        model.transformer.wte.output = embeddings

        output = model.generator.output.save()

print(model.tokenizer.decode(output[0]))
print(model.tokenizer.decode(output[1]))

这会产生:

麦迪逊广场花园位于纽约市。
_ _ _ _ _ _ _ _ _ _纽约市。

我们也可以输入预先保存的嵌入张量,如下所示:

from nnsight import LanguageModel

model = LanguageModel('openai-community/gpt2', device_map='cuda')

with model.generate(max_new_tokens=3) as tracer:
    
    with tracer.invoke("麦迪逊广场花园位于纽约") as invoker:

        embeddings = model.transformer.wte.output.save()

with model.generate(max_new_tokens=3) as tracer:

    with tracer.invoke("_ _ _ _ _ _ _ _ _ _") as invoker:

        model.transformer.wte.output = embeddings.value
临时模块

我们还可以在计算过程中的任何时候应用模型模块树中的模块,即使它是无序的。

from nnsight import LanguageModel
import torch

model = LanguageModel("openai-community/gpt2", device_map='cuda')

with model.generate('埃菲尔铁塔位于') as generator:

  hidden_states = model.transformer.h[-1].output[0]
  hidden_states = model.lm_head(model.transformer.ln_f(hidden_states)).save()
  tokens = torch.softmax(hidden_states, dim=2).argmax(dim=2).save()
        
print(hidden_states)
print(tokens)
print(model.tokenizer.decode(tokens[0]))

这里我们像往常一样获取最后一层的隐藏状态。我们还链式应用model.transformer.ln_f和model.lm_head以便将隐藏状态"解码"到词汇空间。 应用softmax然后argmax允许我们将词汇空间的隐藏状态转换为实际的令牌,然后我们可以使用分词器来解码。

输出看起来像:

张量([[[   -36.2874,   -35.0114,   -38.0793,  ...,   -40.5163,   -41.3759,
            -34.9193],
          [  -68.8886,   -70.1562,   -71.8408,  ...,   -80.4195,   -78.2552,
            -71.1206],
          [  -82.2950,   -81.6519,   -83.9941,  ...,   -94.4878,   -94.5194,
            -85.6998],
          ...,
          [-113.8675,  -111.8628,  -113.6634,  ...,  -116.7652,  -114.8267,
           -112.3621],
          [  -81.8531,   -83.3006,   -91.8192,  ...,   -92.9943,   -89.8382,
            -85.6898],
          [-103.9307,  -102.5054,  -105.1563,  ...,  -109.3099,  -110.4195,
           -103.1395]]], 设备='cuda:0')
张量([[  198,    12,   417,  8765,   318,   257,   262,  3504,  7372,  6342]],
       设备='cuda:0')

-埃菲尔铁塔是巴黎市中心

---

更多示例可以在[nnsight.net](https://www.nnsight.net)找到

### 引用

如果您在研究中使用`nnsight`，请使用以下方式进行引用

```bibtex
@article{fiottokaufman2024nnsightndifdemocratizingaccess,
      title={NNsight和NDIF：实现基础模型内部的民主化访问}, 
      author={Jaden Fiotto-Kaufman and Alexander R Loftus and Eric Todd and Jannik Brinkmann and Caden Juang and Koyena Pal and Can Rager and Aaron Mueller and Samuel Marks and Arnab Sen Sharma and Francesca Lucchetti and Michael Ripa and Adam Belfki and Nikhil Prakash and Sumeet Multani and Carla Brodley and Arjun Guha and Jonathan Bell and Byron Wallace and David Bau},
      year={2024},
      eprint={2407.14561},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2407.14561}, 
}
相关项目
项目侧边栏1项目侧边栏2
推荐项目
Project Cover

豆包MarsCode

豆包 MarsCode 是一款革命性的编程助手，通过AI技术提供代码补全、单测生成、代码解释和智能问答等功能，支持100+编程语言，与主流编辑器无缝集成，显著提升开发效率和代码质量。

Project Cover

AI写歌

Suno AI是一个革命性的AI音乐创作平台，能在短短30秒内帮助用户创作出一首完整的歌曲。无论是寻找创作灵感还是需要快速制作音乐，Suno AI都是音乐爱好者和专业人士的理想选择。

Project Cover

有言AI

有言平台提供一站式AIGC视频创作解决方案，通过智能技术简化视频制作流程。无论是企业宣传还是个人分享，有言都能帮助用户快速、轻松地制作出专业级别的视频内容。

Project Cover

Kimi

Kimi AI助手提供多语言对话支持，能够阅读和理解用户上传的文件内容，解析网页信息，并结合搜索结果为用户提供详尽的答案。无论是日常咨询还是专业问题，Kimi都能以友好、专业的方式提供帮助。

Project Cover

阿里绘蛙

绘蛙是阿里巴巴集团推出的革命性AI电商营销平台。利用尖端人工智能技术，为商家提供一键生成商品图和营销文案的服务，显著提升内容创作效率和营销效果。适用于淘宝、天猫等电商平台，让商品第一时间被种草。

Project Cover

吐司

探索Tensor.Art平台的独特AI模型，免费访问各种图像生成与AI训练工具，从Stable Diffusion等基础模型开始，轻松实现创新图像生成。体验前沿的AI技术，推动个人和企业的创新发展。

Project Cover

SubCat字幕猫

SubCat字幕猫APP是一款创新的视频播放器，它将改变您观看视频的方式！SubCat结合了先进的人工智能技术，为您提供即时视频字幕翻译，无论是本地视频还是网络流媒体，让您轻松享受各种语言的内容。

Project Cover

美间AI

美间AI创意设计平台，利用前沿AI技术，为设计师和营销人员提供一站式设计解决方案。从智能海报到3D效果图，再到文案生成，美间让创意设计更简单、更高效。

Project Cover

AIWritePaper论文写作

AIWritePaper论文写作是一站式AI论文写作辅助工具，简化了选题、文献检索至论文撰写的整个过程。通过简单设定，平台可快速生成高质量论文大纲和全文，配合图表、参考文献等一应俱全，同时提供开题报告和答辩PPT等增值服务，保障数据安全，有效提升写作效率和论文质量。

使用协议隐私政策广告服务
投诉举报邮箱: service@vectorlightyear.com
@2024 懂AI·鲁ICP备2024100362号-6·鲁公网安备37021002001498号