NER项目 P8 CRF维特比解码和模型训练

前文已经定义好了 BiLSTM、Linear 层，最后还需要加一层 CRF，再结合维特比解码，输出目标值。

另外一个重要的步骤，就是定义损失函数，CRF 解码后的损失计算，不同于分类任务的交叉熵损失，还需要考虑到顺序的问题。

定义好以上内容后，就可以进行模型训练了。

Pytorch-crf文档：https://pytorch-crf.readthedocs.io/

模型结构

代码示例

1、CRF

三方包有改动，有微调~

# model.py
class Model(nn.Module):
    def __init__(self):
        ......
        self.crf = CRF(TARGET_SIZE)

    def forward(self, input, mask):
        ......
        return self.crf.decode(out, mask)

    def loss_fn(self, input, target, mask):
        y_pred = self._get_lstm_feature(input)
        return -self.crf.forward(y_pred, target, mask, reduction='mean')

2、模型训练

from utils import *
from model import *
from config import *

if __name__ == '__main__':
    dataset = Dataset()
    loader = data.DataLoader(
        dataset,
        batch_size=100,
        shuffle=True,
        collate_fn=collate_fn,
    )

    model = Model()
    optimizer = torch.optim.Adam(model.parameters(), lr=LR)

    for e in range(EPOCH):
        for b, (input, target, mask) in enumerate(loader):
            y_pred = model(input, mask)
            loss = model.loss_fn(input, target, mask)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            if b % 10 == 0:
                print('>> epoch:', e, 'loss:', loss.item())

        torch.save(model, MODEL_DIR + f'model_{e}.pth')

模型整体的计算量还是挺大的，一般是要用GPU跑的，但 Mac 上调不了 GPU，盲改容易报错，会给大家学习带来困扰，就索性不改了。

大家有 GPU 的设备的，可以自己尝试改一下，只需要记住一点，就是模型和数据保持一致就好。