NLP入门——卷积语言模型的搭建、训练与预测

语言模型建模是针对句子建模，主要分为掩码语言模型和自回归语言模型。
我们从corpus中截取一句话作为例子，这句话是bpe分词后的句子：

1994 年 5 月 17 日 安全 理事会 第 33@@ 77 次 会议 通过

掩码语言模型的主要机制是：例如将33@@ 77这两个分词做掩码，这样句子变成了

1994 年 5 月 17 日 安全 理事会 第 [MASK] [MASK] 会议 通过

将带有掩码的句子输入神经网络，让模型预测掩码位置的正确分词，例如上文正确的分词为 33@@ 、77.在分词的词表中，训练的目的是为了使得预测这句话的mask时，33@@、77分词的概率最高，从而选取这两个分词作为正确的预测结果。

自回归语言模型的机制是：给出这句话的前i-1个词，来预测第i个词；随后加入第i个词，预测这句话的第i+1个词。例如：

1994 -> 年
1994 年 -> 5
1994 年 5 -> 月
…
1994 年 5 月 17 日 安全 理事会 第 33@@ 77 次 会议 -> 通过

在自回归语言模型主要有两个问题：
1.句首的词怎么来，比如本句话的1994.
2.句末的分词“通过”生成之后，是继续生成还是结束？
因此，对句子的处理我们要加两个特殊标记：句首加[SOS]（start of sentence），句末加[EOS]（end of sentence）.例如：

[SOS] 1994 年 5 月 17 日 安全 理事会 第 33@@ 77 次 会议 -> 通过 [EOS]

[SOS]作触发器，当遇到时开始了第一个词的生成以及整个句子的生成；[EOS]是结束的作用，当分类产生的概率[EOS]最大的时候，就认为不需要继续生成了。

模型的搭建

首先我们对分词后的语料按句长升序排序，设置最长句长为254，这样加上 eos 和 sos 后最大句长为256。

#sorti.py
#encoding: utf-8
import sys

def handle(srcf, srts, max_len=256): #设置句长最大256，大于则丢弃

    # {length: {src}}  外层dict，内层set
    _maxlen = max_len - 2   #减去[SOS]和[EOS]
    data = {}
    with open(srcf, "rb") as fsrc:
        for ls in fsrc:
            ls = ls.strip()
            if ls:
                ls = ls.decode("utf-8")
                _ = len(ls.split()) #获取句子的分词个数
                if _ <= max_len:
                    if _ in data:   #若已有这个长度在data中
                        if ls not in data[_]: #去重，重复的跳过
                            data[_].add(ls)   #不重复的添加
                    else:
                        data[_] = set([ls])   #转化成set去重
    ens = "\n".encode("utf-8")
    with open(srts, "wb") as fsrc: #写入
        for _l in sorted(data.keys()):  #按照句子长度从小到大排
            lset = data[_l] #取出句长对应的set
            fsrc.write("\n".join(lset).encode("utf-8")) #在每个句子间插入换行符
            fsrc.write(ens) #每个句子后插入换行

if __name__ == "__main__":
    handle(*sys.argv[1:3])

在命令行执行并查看结果文件：

:~/nlp/lm$ python sorti.py ../token/zh.bpe zh.srt
:~/nlp/lm$ less zh.srt

由于自回归语言模型的机制，我们通过已有的分词预测句子中下一个分词，则不断将预测的分词加入句子来预测下一个分词，那么输入就是变长的。为了解决这个问题，我们使用卷积操作，设置每次固定的上下文长度ctx_len，来预测下一个分词。
假设ctx_len=4，如果我们有序列0 1 2 3 4 5 6 7 8，则有：
0 1 2 3 -> 4
1 2 3 4 -> 5
2 3 4 5 -> 6
3 4 5 6 -> 7
4 5 6 7 -> 8
因此，若每个batch中句子的长度为seql，则我们每次需要取seql - ctx_len个长度的分词，共取ctx_len次。

>>> import torch        
>>> a = torch.arange(9,dtype=torch.long)
>>> ctx_len = 4
>>> nlen = 9 - ctx_len
>>> _ = [a.narrow(0,i,nlen) for i  in range(ctx_len)]
>>> _
[tensor([0, 1, 2, 3, 4]), tensor([1, 2, 3, 4, 5]), tensor([2, 3, 4, 5, 6]), tensor([3, 4, 5, 6, 7])]
>>> b = torch.stack(_,dim=-1)
>>> b
tensor([[0, 1, 2, 3],
        [1, 2, 3, 4],
        [2, 3, 4, 5],
        [3, 4, 5, 6],
        [4, 5, 6, 7]])

#CNNLM.py
#encoding: utf-8
import torch
from torch import nn

class NNLayer(nn.Module):
    def __init__(self, isize, hsize, dropout,norm_residual=True,
    **kwargs):
        super(NNLayer, self,).__init__()   ##调用父类的初始化函数
        self.net = nn.Sequential(nn.Linear(isize, hsize),
        nn.ReLU(inplace=True), #设置relu激活函数，inplace=True在原始张量上进行
        nn.Dropout(p=dropout, inplace=False),#设置丢弃率防止过拟合，同时创建一个新的张量
        nn.Linear(hsize, isize, bias=False), nn.Dropout(p=dropout, inplace=True))
        self.normer = nn.LayerNorm(isize) #做归一化
        self.norm_residual = norm_residual #设置变量存储做判断
        
    def forward(self, input):
        
        _ = self.normer(input) #稳定之后的结果 
        return (_ if self.norm_residual else input) + self.net(_)
        #如果参数初始化做的好，就用LayerNorm后的值，否则用原始值

class NNLM(nn.Module):
    
    def __init__(self, vcb_size, isize, hsize, dropout,
    nlayer, ctx_len=4, bindemb=True, **kwargs): #有多少个词，就分多少类，类别数为vcb_size
        super(NNLM, self).__init__()
        self.emb = nn.Embedding(vcb_size, isize,
        padding_idx=0)                #<pad>的索引为0
        self.comp = nn.Linear(ctx_len * isize, isize, bias=False) #将4个词的向量降维为isize
        self.drop = nn.Dropout(p=dropout, inplace=True) #embedding后dropout
        self.nets = nn.Sequential(*[NNLayer(isize, hsize, dropout)
        for _ in range(nlayer)])
        self.classifier = nn.Linear(isize, vcb_size)
        if bindemb:
            self.classifier.weight = self.emb.weight#将emb的权重赋给分类器
        self.normer = nn.LayerNorm(isize)
        self.out_normer = nn.LayerNorm(isize)
        self.ctx_len = ctx_len      #保存上下文长度
        
    # input: (bsize, seql-1) 句数，句长-1 由于最后一个词是预测不作为输入
    def forward(self, input):
        
        out = self.emb(input)
        # out: (bsize, seql-1, isize)
        _nlen = input.size(1) - self.ctx_len + 1# 每次取_nlen长度的分词数
        _ = []
        for i in range(self.ctx_len):
            _.append(out.narrow(1, i, _nlen)) #从seql这一维取，每次从i开始取，取_nlen长度
        
        out = torch.cat(_,dim=-1)    
        #[(bsize, _nlen, isize)] -> (bsize, _nlen, ctx_len * isize)  
        
        out = self.normer(self.drop(out))
        out = self.comp(out)  #降维
        #(bsize, _nlen, ctx_len * isize) -> (bsize, _nlen, isize)   
        out = self.normer(out) #使用归一化，使模长均匀 
        out = self.out_normer(self.nets(out))
        return self.classifier(out) #分类产生参数

分词词典映射

#vcb.py
#encoding: utf-8
import sys

#预定义规则
init_vcb ={"<pad>": 0, "<sos>": 1, "<eos>": 2} #将占位符、开始标志和结束标志的ID映射为0、1、2
init_token_id = len(init_vcb)      #设置其他词的ID从3开始编号

def get_vcb(*files): #词频统计
    
    #{word: freq}
    vcb = {}
    for file in files:
        with open(file, "rb") as f:
            for line in f:
                line = line.strip()
                if line:
                    for word in line.decode("utf-8").split():
                        vcb[word] = vcb.get(word,0) + 1
                        
    return vcb

# freq word0 word1 word2 将同一频率的词放在一起
def save_vcb(vcb, file):
    
    # {freq: [word]}
    rs = {}
    for word, freq in vcb.items():
        if freq in rs: #若此词频项已存在，则添加该词到list
            rs[freq].append(word)
        else:          #若此项不存在，则将该词作为list第一个词
            rs[freq] = [word]

    with open(file, "wb") as f:
        for freq in sorted(rs.keys(), reverse=True):#从大到小读rs的频率，reverse=True说明按照从大到小排序
            f.write(("%d\t%s\n"%(freq,"\t".join(rs[freq]))).#将同一词频的词合成一个字符串，词间用制表符分开
            encode("utf-8"))
            
def load_vcb(file, vanilla = False, init_vcb = init_vcb, init_token_id = init_token_id):#为每个分词建立索引
    #vanilla变量是为了区分是否使用了<pad>，若使用则编号从1开始否则从0开始
    vcb, cid = ({}, 0) if vanilla else (init_vcb.copy(), 
    init_token_id)
    
    with open(file, "rb") as f:
        for line in f:
            line = line.strip()
            if line:
                for word in line.decode("utf-8").split()[1:]: #舍弃词频，只需要保留分词
                    vcb[word] = cid
                    cid += 1
    #为每个分词建立索引
    return vcb
      
reverse_vcb = lambda vcb: {v: k for k, v in vcb.items()}
#使用匿名函数将 词->映射 转换为 映射->词

if __name__ == "__main__":
    save_vcb(get_vcb(*sys.argv[1:-1]), sys.argv[-1])
    #最后一个参数为存词频————分词的文件

我们对已经排好序的句子进行词频的收集，在命令行执行：

:~/nlp/lm$ python vcb.py zh.srt zh.vcb
:~/nlp/lm$ less zh.vcb

我们可以看到收集好的词典按照词频从高到低保存了所有分词。

文本转张量

#mkh5.py
#encoding: utf-8
import sys
from h5py import File as h5File
import numpy
from vcb import load_vcb
#导入加载词典的函数

def batch_loader(fsrc, max_tokens = 4096, **kwargs):#返回一批一batch的数据，设置每个batch最多存放2048个子词 
    ri = []
    mlen = n = 0    #n记录当前收集了多少条句子，mlen记录当前收集的句子长度
    with open(fsrc, "rb") as fs:
        for ls in fs:
            ls = ls.strip() 
            if ls:
                ls =ls.decode("utf-8").split()
                _l = len(ls)            #当前行中的分词个数
                _mlen = max(_l, mlen)   #当前行或当前batch中句子的长度
                _n = n + 1
                if (_n * _mlen) > max_tokens: #如果把添加了这句话的 句数*分词数量 大于最大值则不能放
                    if ri:       #如果ri不为空
                        yield ri, mlen  #返回ri和原来的句子长度
                    ri, mlen, n = [ls], _l, 1 #返回后重新初始化，将本句加入新的batch
                else: #如果不超过当前长度，则将此句添加到batch中
                    ri.append(ls)
                    mlen, n = _mlen, _n #更新句子长度与句子数量
    if ri: #最后若仍然有数据，则返回为一个新的batch
        yield ri, mlen
                
def map_instance(lin, vcb, sos_id=1, eos_id=2):
    
    rs = [sos_id]  #添加开始标志
    rs.extend([vcb[_word] for _word in lin if _word  in vcb])
    rs.append(eos_id) #添加结束标志

    return rs

def batch_mapper(fsrc, vcbs, **kwargs): #将分词变索引
    for ri, mlen in batch_loader(fsrc, **kwargs):
        yield [map_instance(_s, vcbs) for _s in ri], mlen + 2
    #遍历每个batch中的句子，返回每个batch中每个分词的索引，以及batch长度
      
def pad_batch(lin, mlen, pad_id = 0):#补<pad>的函数

    rs = []
    for lu in lin: #每个batch中的每句
        _d = mlen - len(lu) #当前此句需要补<pad>的个数
        if _d > 0:
            lu.extend([pad_id for _ in range(_d)])#extend函数用来拼接两个列表。补_d个<pad>的索引0
        rs.append(lu)
        
    return rs #返回的是均已对齐的每个batch
    
def batch_padder(fsrc, vcbs, **kwargs):
    for ri, mlen in batch_mapper(fsrc, vcbs, **kwargs):
        yield pad_batch(ri, mlen) #返回的是每个已补齐的batch，以及batch中句子的长度
        
def handle(fsrc, fvcbs, frs, **kwargs):
    vcbs = load_vcb(fvcbs, vanilla = False)
    with h5File(frs, "w", libver = 'latest', track_order = False) as h5f:#libver使用最新的，track_order表示无需记录顺序
        src_grp = h5f.create_group("src", track_order=False)    #创建两个组，分别放句子和标签 
        for i, ri in enumerate(batch_padder(fsrc, vcbs, **kwargs)):
            ri = numpy.array(ri, dtype = numpy.int32) #转化成numpy数组并设置数据类型，target的数据很小，所以我们只需要int16存储
            src_grp.create_dataset(str(i), data=ri, compression="gzip",
            compression_opts=9, shuffle=True ) #设置压缩存储节省空间，压缩等级设置为最大压缩代价9
            
        h5f["nword"] = numpy.array([len(vcbs)], dtype=numpy.int32) #存储总词数
        h5f["ndata"] = numpy.array([i + 1], dtype=numpy.int32) #存储总batch数
   
if __name__ == "__main__":
    handle(*sys.argv[1:4])

在命令行执行：

:~/nlp/lm$ python mkh5.py zh.srt zh.vcb train.h5
:~/nlp/lm$ h5ls -d train.h5 
ndata                    Dataset {1}
    Data:
         83685
nword                    Dataset {1}
    Data:
         42585
src                      Group

可以看到我们总的训练集张量共有42585个分词以及83685个batch。

模型的训练

#train.py
#encoding: utf-8

import torch
from torch import nn
from CNNLM import NNLM #导入模型
from h5py import File as h5File #读训练数据
from math import sqrt
from random import shuffle #使输入数据乱序，使模型更均衡
from lrsch import SqrtDecayLR
from tqdm import tqdm

train_data = "train.h5"#之前已经张量转文本的h5文件
isize = 64              
hsize = isize * 2       #设置初始参数
dropout = 0.3           #设置丢弃率
nlayer = 4              #设置层数
gpu_id = -1             #设置是否使用gpu
lr = 1e-3               #设置初始学习率
max_run = 8           #设置训练轮数
#early_stop = 16         #设置早停轮数
ctx_len = 4             #设置上下文长度
nreport = 5000          #每训练5000个batch打印一次
tokens_optm = 25000     #设置更新参数的词数阈值

def init_model_parameters(modin): #初始化模型参数
    with torch.no_grad():         #不是训练不用求导
        for para in modin.parameters():
            if para.dim() > 1:          #若维度大于1，说明是权重参数
                _ = 1.0 / sqrt(para.size(-1))
                para.uniform_(-_,_)     #均匀分布初始化
        for _m in modin.modules():      #遍历所有小模型
            if isinstance(_m, nn.Linear):#如果小模型是linear类型
                if _m.bias is not None: #初始化bias
                    _m.bias.zero_()
                elif isinstance(_m, nn.LayerNorm):#初始化LayerNorm参数
                    _m.weight.fill_(1.0)
                    _m.bias.zero_()
    return modin

def train(train_data, tl, model, lossf, optm, cuda_device,
 nreport=nreport, tokens_optm=tokens_optm, ctx_len=ctx_len):#nreport为每训练一部分词打一次epoch
    
    model.train() #设置模型在训练的模式
    src_grp = train_data["src"] #从输入数据中取出句子
    _l = 0.0  #_l用来存当前loss
    _t = 0    #_t用来存句子数
    _lb = 0.0
    _tb = 0
    _tom = 0
    _min_len_thres = ctx_len + 1
    for _i, _id in tqdm(enumerate(tl, 1)):
        seq_batch = torch.from_numpy(src_grp[_id][()])
        #seq_batch:[bsize, seql]
        _seqlen = seq_batch.size(-1)  #取出每个batch的句长
        if _seqlen < _min_len_thres:  #至少有ctx_len个词才能运行，否则跳过
            continue
        if cuda_device is not None:
            seq_batch = seq_batch.to(cuda_device, non_blocking=True)
             #将数据放在同一gpu上
        seq_batch = seq_batch.long()   #数据转换为long类型
        seq_i = seq_batch.narrow(1, 0, _seqlen - 1) #训练数据读取bsize个seql-1的数据
        #seq_i:[bsize, seql-1]
        seq_o = seq_batch.narrow(1, ctx_len, _seqlen - ctx_len) #预测数据读取ctx_len~seql-1的预测数据
        #seq_o:[bsize, seql-ctx_len]
        
        #out: {bsize, seql-ctx_len, vcb_size} vcb_size即预测类别数
        
        out = model(seq_i)          #获得模型结果
        loss = lossf(out.view(-1, out.size(-1)), seq_o.contiguous().view(-1)) 
        #转换out维度为[bsize*(seql-ctx_len),vcb_size],seq_o:[bsize*(seql-ctx_len)]
        _lossv = loss.item()
        _l += _lossv        #整个训练集的loss
        _lb += _lossv       #每个batch的loss
        _n = seq_o.ne(0).int().sum().item() #seq_o中不是<pad>的位置的数量
        _t += _n            #整个训练集的分词数
        _tb += _n           #每个batch的分词数
        _tom += _n
        loss.backward()                 #反向传播求导
        if _tom > tokens_optm:          #当词数大于时更新参数
            optm.step()                     #参数的更新
            optm.zero_grad(set_to_none=True)#参数更新后清空梯度
            _tom = 0
        if _i % nreport == 0:   #每训练5000个batch打印一次
            print("Average loss over %d tokens: %.2f"%(_tb, _lb/_tb))
            _lb = 0.0
            _tb = 0
            save_model(model, "checkpoint.pt") #暂存检查点模型
                    
    return _l / _t #返回总的loss

def save_model(modin, fname): #保存模型所有内容 权重、偏移、优化
    
    torch.save({name: para.cpu() for name, para in
    model.named_parameters()}, fname)

t_data = h5File(train_data, "r")#以读的方式打开训练数据

vcb_size = t_data["nword"][()].tolist()[0] #将返回的numpy的ndarray转为list
#在我们的h5文件中存储了nword： 总词数

model = NNLM(vcb_size, isize, hsize, dropout, nlayer, ctx_len=ctx_len)
model = init_model_parameters(model) #在cpu上初始化模型
lossf = nn.CrossEntropyLoss(reduction='sum', ignore_index=0,
label_smoothing=0.1)
#设置ignore_index=0，即忽略<pad>的影响

if (gpu_id >= 0) and torch.cuda.is_available(): #如果使用gpu且设备支持cuda
    cuda_device = torch.device("cuda", gpu_id)  #配置gpu
    torch.set_default_device(cuda_device)
else:
    cuda_device = None

if cuda_device is not None:                     #如果要用gpu
    model.to(cuda_device)                       #将模型和损失函数放在gpu上
    lossf.to(cuda_device)

optm = torch.optim.Adam(model.parameters(), lr=lr, 
betas=(0.9, 0.98), eps=1e-08)
#使用model.parameters()返回模型所有参数，
lrm = SqrtDecayLR(optm, lr) #将优化器和初始学习率传入

tl = [str(_) for _ in range(t_data["ndata"][()].item())] #获得字符串构成的训练数据的list

for i in range(1, max_run + 1):
    shuffle(tl)         #使数据乱序
    _tloss = train(t_data, tl, model, lossf, optm,
    cuda_device)  #获取每轮训练的损失
    print("Epoch %d: train loss %.2f"%(i, _tloss)) #打印日志
    save_model(model, "eva.pt")
    lrm.step() #每轮训练后更新学习率
    
t_data.close()

在训练脚本中，我们要注意训练集很大，所以设置每nreport 个batch打印一次结果。更新参数时也设置一个阈值来多轮更新一次而不是每个batch更新一次。
由于整体训练模型开销大，因此我们每nreport个batch暂存一次检查点模型方便解码和后续工作的推进。
我们训练的过程就是将句子的前seql-1个词输入模型（seql为句子长度，即句子分词个数），将训练集句子第ctx_len以后的词作为标签（模型从前ctx_len个词预测第ctx_len+1的词，再添加再预测以此类推），直到下标为 seql-ctx_len-2 到 seql-2 预测下标为 seql-1 的词。至此句长为seql的0~seql-1的词都过了训练。计算预测分词和标签的交叉熵损失。

模型的预测

预测数据准备

我们选取两句话让模型预测之后的句子：

审议 了 1994 年 5 月 13 日 秘书长 的 报告 ( S / 1994 / 565 ) ,
重申 其 1993 年 9 月 29 日 关于 联合国 行动 安全 的 第 868 ( 1993 ) 号 决议 ,

我们做测试用的句子需满足句长（即分词个数）大于5，即大于ctx_len+1.
我们将其存入文件中，然后保存为h5文件：

:~/nlp/lm$ less test.bpe 
审议 了 1994 年 5 月 13 日 秘书长 的 报告 ( S / 1994 / 565 ) ,
重申 其 1993 年 9 月 29 日 关于 联合国 行动 安全 的 第 868 ( 1993 ) 号 决议 ,

由于这两句是长度升序排列，所以省去排序的步骤。

:~/nlp/lm$ python mkh5.py test.bpe zh.vcb test.h5 
:~/nlp/lm$ h5ls -d test.h5/src
0                        Dataset {2, 23}
    Data:
         1, 159, 14, 1246, 12, 112, 20, 352, 35, 125, 3, 34, 11, 532, 13, 1246, 13, 12210, 10, 4, 2, 0,
         0, 1, 489, 39, 1459, 12, 165, 20, 600, 35, 41, 33, 78, 79, 3, 16, 20380, 11, 1459, 10, 63, 67,
         4, 2

我们可以看到第一句话的张量是：[ 1, 159, 14, 1246, 12, 112, 20, 352, 35, 125, 3, 34, 11, 532, 13, 1246, 13, 12210, 10, 4, 2, 0, 0,]
第二句话的张量是：[1, 489, 39, 1459, 12, 165, 20, 600, 35, 41, 33, 78, 79, 3, 16, 20380, 11, 1459, 10, 63, 67, 4, 2]

他们都满足：开头是<sos>，索引为1；结尾是<eos>，索引为2。第一句话由于要和第二句长度对齐构成一个batch，在末尾补了两个<pad>，索引为0.
而我们预测续写句子时，在测试集我们首先不需要结尾标记2以及为补齐句子长度而在末尾添加的0.于是我们将pad占位符放在句子的开头，对mkh5.py脚本做以下修改：

#mkh5test.py
def pad_batch(lin, mlen, pad_id = 0):#补<pad>的函数

    rs = []
    for lu in lin: #每个batch中的每句
        _d = mlen - len(lu) #当前此句需要补<pad>的个数
        if _d > 0:
            _ = [pad_id for _ in range(_d)]
            #lu.extend([pad_id for _ in range(_d)])#extend函数用来拼接两个列表。补_d个<pad>
            _.extend(lu) #将原句放在<pad>后
        else:
            _ = lu
        rs.append(_)
        
    return rs #返回的是均已对齐的每个batch
def map_instance(lin, vcb, sos_id=1, eos_id=2):
    
    rs = [sos_id]
    rs.extend([vcb[_word] for _word in lin if _word  in vcb])
    # rs.append(eos_id) 测试时句末无需加eos

    return rs

再次保存h5文件后：

:~/nlp/lm$ python mktesth5.py test.bpe zh.vcb test.h5 
:~/nlp/lm$ h5ls -d test.h5/src
0                        Dataset {2, 22}
    Data:
         0, 0, 1, 159, 14, 1246, 12, 112, 20, 352, 35, 125, 3, 34, 11, 532, 13, 1246, 13, 12210, 10, 4,
         1, 489, 39, 1459, 12, 165, 20, 600, 35, 41, 33, 78, 79, 3, 16, 20380, 11, 1459, 10, 63, 67, 4

第一句话张量:[ 0, 0, 1, 159, 14, 1246, 12, 112, 20, 352, 35, 125, 3, 34, 11, 532, 13, 1246, 13, 12210, 10, 4,]
第二句话张量：[ 1, 489, 39, 1459, 12, 165, 20, 600, 35, 41, 33, 78, 79, 3, 16, 20380, 11, 1459, 10, 63, 67, 4]

解码与预测

模型的解码

    #CNNLM.py
    # input: (bsize, seql)
    def decode(self, input, maxlen=50):  #设置最多生成50个分词
        
        rs = input
        bsize =input.size(0)
        done_trans = None   #记录是否完成生成
        for i in range(maxlen):
            out = rs.narrow(1, rs.size(1) - self.ctx_len, self.ctx_len) #利用后四个词预测下一个词
            # (bsize, ctx_len, isize)
            out = self.emb(out).view(bsize, -1) 
            # (bsize, ctx_len*isize)
            out = self.normer(self.comp(self.drop(out)))
            out = self.out_normer(self.nets(out))
            out = self.classifier(out).argmax(-1,keepdim=True) #取最后一维分数最高的索引
            #这一步对应分类，keepdim=True保持维度便于拼接
            # out:(bsize, vcb_size) -> (bsize, 1)
            rs = torch.cat([rs, out], dim=1)    #将预测的词拼接到原句后，在第一维度即seql后
            _eos = out.eq(2) #当遇到<eos>解码停止
            # _eos:(bsize, 1)
            if done_trans is None:
                done_trans = _eos
            else:
                done_trans |= _eos #将_eos中的True赋给done_trans
            #_eos中的元素如果为True，则说明在该索引位置上out值为2即结束标志
            if done_trans.all().item(): #当全都为True，说明此batch中所有句子预测都为<eos>，即解码完成
                break
      
        return rs

首先，我们需要在模型文件的NNLM类中添加decode方法用于解码，decode方法将输入句子的最后ctx_len个分词作为模型的输入，预测下一个分词后将其拼接到句子中，下次再取后三个以及新生成的分词作为输入预测；
重复此步骤直至生成超过50个分词或者这一个batch中生成全部为<eos>结束标志，则停止。
self.classifier(out).argmax(-1,keepdim=True) 这句话在执行argmax函数前，classifier分类器分类后的size是(bsize, vcb_size)，所以执行argmax函数在最后一维取的下标索引即为分词在词表中的下标。

模型的预测

#predict.py
#encoding: utf-8

import sys
import torch
from CNNLM import NNLM               #读模型
from h5py import File as h5File         #读文件
from vcb import load_vcb, reverse_vcb   #获取词表

isize = 64              
hsize = isize * 2       #设置初始参数
dropout = 0.3           #设置丢弃率
nlayer = 4              #设置层数
gpu_id = -1             #设置是否使用gpu
ctx_len = 4             #设置上下文长度

def extract(lin, vcb): #提取结果的函数
    
    rs = []
    for lu in lin:
        if lu > 1:
            if lu == 2:
                break
            else:
                rs.append(vcb[lu])  #返回索引对应词典中的分词
    
    return " ".join(rs)     #返回空格分隔的解码后的字符串

test_data = sys.argv[1]
test_file = h5File(test_data, "r")            #读验证集
vcb_size = test_file["nword"][()].tolist()[0] #获取总词数

tgt_vcb = reverse_vcb(load_vcb(sys.argv[2], vanilla=False))
#vanilla设置为false，读取词表时需考虑到特殊标记：0,1,2

model = NNLM(vcb_size, isize, hsize, dropout, nlayer, ctx_len=ctx_len)
model_file = sys.argv[-1]       #获取模型
with torch.no_grad():           #避免求导
    _ = torch.load(model_file)  #加载词典
    for name, para in model.named_parameters():
        if name in _:
            para.copy_(_[name]) #从词典里取出name的参数

if (gpu_id >= 0) and torch.cuda.is_available():
    cuda_device = torch.device("cuda", gpu_id)
    torch.set_default_device(cuda_device)
else:
    cuda_device = None 
if cuda_device is not None:
    model.to(cuda_device)       #判断是否使用cuda

model.eval()

src_grp = test_file["src"]
ens = "\n".encode("utf-8")
with torch.no_grad(), open(sys.argv[3],"wb") as f: #解码避免求导，将预测标签按行写入文件
    for _ in range(test_file["ndata"][()].item()):#每个batch上遍历
        seq_batch = torch.from_numpy(src_grp[str(_)][()])
        if cuda_device is not None:
            seq_batch = seq_batch.to(cuda_device, non_blocking=True)
        seq_batch = seq_batch.long()    #s数据类型转换
        output = model.decode(seq_batch).tolist() #将解码后的numpy转为list
        output = [extract(_, tgt_vcb) for _ in output] #将张量转为文本
        f.write("\n".join(output).encode("utf-8"))
        f.write(ens)                #每个batch间还应有换行
 
test_file.close()

在命令行执行：

:~/nlp/lm$ python predict.py test.h5 zh.vcb pred.txt checkpoint.pt 
:~/nlp/lm$ less pred.txt

可以看到模型对测试句子做了续写。