手把手教你用PyTorch搭一个能聊天的AI机器人，别被那些高级框架唬住了

说实话，我最早接触聊天机器人那会儿，脑子里全是各种高大上的框架——什么Transformer、GPT、BERT，光是看名字就觉得这东西离自己很远，直到有一天，我实在闲得慌，翻了翻PyTorch的官方文档，又对照着网上一些零散的教程，折腾了一个周末，居然真搞出来一个能跟我瞎扯几句的聊天机器人，虽然它经常答非所问，但那种“我自己写出来的”成就感,比用现成API爽太多了。

今天我就把当时踩过的坑和思路写下来，不是什么专业教程，就是跟你聊聊怎么从零开始用PyTorch搭一个能聊天的东西，别怕,代码其实没你想得那么玄乎。

先别想太复杂，我们到底在干什么？

很多人一听到“聊天机器人”，第一反应就是ChatGPT那种级别的，兄弟，醒醒，那是几百亿参数、几万张显卡堆出来的东西，我们要做的，说白了就是让机器学会“根据你说的话，猜我该回什么”。

最朴素的做法是什么？sequence-to-sequence模型，就是你把一句话扔进去，模型把它编码成一个向量，然后解码器从这个向量里“翻译”成回复，跟机器翻译本质上是一个逻辑,只不过源语言和目标语言都变成了中文。

我当时选的是LSTM，因为RNN虽然简单但梯度消失太厉害，LSTM好歹能记住点上下文，当然你也可以用GRU，省显存，不过现在回头看，其实用Transformer的编码器层更稳，但那时我还没学会怎么调注意力机制,就不给自己加戏了。

数据才是亲爹，写代码反而是最简单的

现在很多教程一上来就贴代码，好像代码写完了机器人就能跟你谈人生一样，扯淡，你拿什么训练？我第一个版本用的是网上随便扒的影视剧台词，结果机器人学了一嘴“你走开”“我恨你”,跟我妈聊天时她差点以为我装了个人工智障。

后来我老老实实去爬了些客服对话数据，整理成“问-答”对，大概弄了两万条，虽然不多，但至少够它学会基本的礼貌了，数据预处理的时候有个坑你得注意——要把所有标点符号统一处理，？”、“？”都转成“？”（中文问号），否则词表会大一倍,而且模型根本分不清它们有什么区别。

对了，分词这块我强烈建议别自己造轮子，直接用jieba或者pyhanlp，省心，词表大小我控制在5000左右，太小了很多词变成UNK，太大了显存撑不住，你别看我这里说得轻巧,当时为了调这个数字我跑了至少七八个版本。

模型结构其实就这几块东西

我用的模型是Encoder-Decoder架构，两边都是双层LSTM，Encoder里把每个词的embedding叠起来，经过LSTM得到一个context vector（其实就是最后一个隐状态），Decoder那边，把上一个预测的词和这个context vector一起送进LSTM，再经过一个全连接层+softmax,选出概率最大的那个词。

听着挺绕是吧？其实代码写起来也就几十行,我贴个核心片段给你感受一下：

class Encoder(nn.Module):
    def __init__(self, vocab_size, embed_size, hidden_size):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_size)
        self.lstm = nn.LSTM(embed_size, hidden_size, batch_first=True)
    def forward(self, x):
        embedded = self.embedding(x)
        output, (hidden, cell) = self.lstm(embedded)
        return hidden, cell

看到没，就这？对，就这，你把词向量维度设成256，隐藏层维度设成512，跑起来问题不大，但如果你显存只有4G，建议把batch_size设小点，我一开始设了64直接爆了,后来降到32才勉强跑通。

Decoder稍微复杂点，因为要逐词生成，不能像Encoder那样一次性全扔进去，但我当时写的时候也没花太多时间，主要是对一下输入输出的维度，别把batch和sequence搞反了就行——这个错我犯了三回才长记性。

训练过程才是最磨人的

你以为模型写好了就能跑？天真，我第一次训练的时候，loss死活不降，代码检查了三遍才发现是embedding层的初始化忘了写，导致词向量全是0，还有一次是teacher forcing比例设得太高，模型完全不会自己生成,只会复制输入。

关于teacher forcing，我的经验是训练刚开始时可以设到0.8甚至0.9，让模型先学会“抄作业”，跑个五六个epoch之后慢慢降到0.2，逼它自己试着生成，当然也有人说直接从头到尾用0.5，我试过,效果真的一般。

学习率我一开始设了0.01，结果loss震荡得像过山车，后来换成0.001加Adam优化器，稳多了，训练了大概20个epoch，花了三个小时,中间还睡了一觉。

别期待太高，但它确实能“聊”

你问最后效果怎么样？老实说，灾难，但也不是不能装装样子，比如我问“你叫什么名字？”，它有时候会回“我是机器人，你是谁？”这种还像那么回事的句子，但更多时候是“你你你你”这样的复读,或者干脆输出一堆UNK。

后来我加了Beam Search（束搜索），把束宽设成3，效果好了一点，至少不会重复到让人头疼，但推理速度慢了一倍,这也是没办法的事。

如果你不想折腾这些细节，其实现在有很多开源的预训练对话模型，比如ChatGLM或者通义千问，直接拿微调一下就行，但话又说回来，亲手从零搭一遍，你对模型内部的理解会深很多——那种感觉像学开车前先拆了一台发动机，虽然麻烦,但再也不怕它坏了。

写在最后

好了，这就是我用PyTorch搭聊天机器人的全过程，不算完整教程，更像是我两个周末的踩坑记录，如果你也想试试，建议从小的词表、小的数据集开始，别一上来就想着搞个大新闻，第一版只要能输出一句人话,就是胜利。

机器人的回复点个赞呗，就当是给我这个折腾了一周的自媒体人一点鼓励，下篇文章我来聊聊怎么用强化学习让聊天机器人学会说人话，不骗你,真的挺有意思的。

（免费申请加入）AI工具导航网

相关标签： # pytorch实现ai聊天机器人