pytorch搭建TextRNN与使用案例

生成训练数据

这里使用随机数生成训练数据，大家在自己写的时候只需要替换这里就OK了：

def get_total_train_data(word_embedding_size, class_count):
    """得到全部的训练数据，这里需要替换成自己的数据"""
    import numpy as np
    x_train = torch.Tensor(np.random.random((1000, 20, word_embedding_size)))  # 维度是 [ 数据量, 一句话的单词个数, 句子的embedding]
    y_train = torch.Tensor(np.random.randint(0, class_count, size=(1000, 1))).long()  # [ 数据量, 句子的分类], 这里的class_count=4，就是四分类任务
    return x_train, y_train

由于多数是用于文本训练，因此这里将word_embedding_size换成自己的embedding，对于分类任务，class_count = N表示N分类任务；回归任务也可以自定义修改代码

构建TextRNN

class TextRnnModel(nn.Module):
    def __init__(self, embedding_size, lstm_hidden_size, output_size, num_layers=2, dropout=0.3):
        super(TextRnnModel, self).__init__()
        self.lstm = nn.LSTM(embedding_size,  # 词嵌入模型词语维度
                            lstm_hidden_size,  # 隐层神经元的维度，为输出的维度
                            num_layers,  # 构建两层的LSTM：堆叠LSTM
                            bidirectional=True,  # 双向的LSTM：词向量从前往后走，再重后往前走，最后拼接起来
                            batch_first=True,  # 把第一个维度的输入作为batch输入的维度
                            dropout=dropout)
        self.fc = nn.Linear(lstm_hidden_size * 2, output_size)
        # nn.Linear(输入维度，输出维度)：[ 上一步输出的LSTM维度*2(双向) , 10分类 ]

    def forward(self, x):
        """前向传播"""
        out, _ = self.lstm(x)  # 过一个LSTM [batch_size, seq_len, embedding]
        out = self.fc(out[:, -1, :])
        return out

开始训练

构建训练数据集

构建batch数据

if __name__ == '__main__':
    epochs = 1000
    batch_size = 30
    embedding_size = 350
    output_class = 14
    x_train, y_train = get_total_train_data(embedding_size, output_class)
    train_loader = Data.DataLoader(
        dataset=Data.TensorDataset(x_train, y_train),  # 封装进Data.TensorDataset()类的数据，可以为任意维度
        batch_size=batch_size,  # 每块的大小
        shuffle=True,  # 要不要打乱数据 (打乱比较好)
        num_workers=6,  # 多进程（multiprocess）来读数据
        drop_last=True,
    )

训练三件套：模型，loss，优化器

	model = TextRnnModel(embedding_size=embedding_size, lstm_hidden_size=200, output_size=output_class)
    cross_loss = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)  # 优化器

开始训练

	model.train()
    for i in range(epochs):
        for seq, labels in train_loader:
            optimizer.zero_grad()
            y_pred = model(seq)  # 压缩维度：得到输出，并将维度为1的去除
            single_loss = cross_loss(y_pred, labels.squeeze())
            # single_loss = cross_loss(y_pred, labels)
            single_loss.backward()
            optimizer.step()
            print("Step: " + str(i) + " loss : " + str(single_loss.detach().numpy()))

完整代码

#!/usr/bin/env Python
# coding=utf-8
import torch
import torch.nn as nn
import torch.utils.data as Data


class TextRnnModel(nn.Module):
    def __init__(self, embedding_size, lstm_hidden_size, output_size, num_layers=2, dropout=0.3):
        super(TextRnnModel, self).__init__()
        self.lstm = nn.LSTM(embedding_size,  # 词嵌入模型词语维度
                            lstm_hidden_size,  # 隐层神经元的维度，为输出的维度
                            num_layers,  # 构建两层的LSTM：堆叠LSTM
                            bidirectional=True,  # 双向的LSTM：词向量从前往后走，再重后往前走，最后拼接起来
                            batch_first=True,  # 把第一个维度的输入作为batch输入的维度
                            dropout=dropout)
        self.fc = nn.Linear(lstm_hidden_size * 2, output_size)
        # nn.Linear(输入维度，输出维度)：[ 上一步输出的LSTM维度*2(双向) , 10分类 ]

    def forward(self, x):
        """前向传播"""
        out, _ = self.lstm(x)  # 过一个LSTM [batch_size, seq_len, embedding]
        out = self.fc(out[:, -1, :])
        return out


def get_total_train_data(word_embedding_size, class_count):
    """得到全部的训练数据，这里需要替换成自己的数据"""
    import numpy as np
    x_train = torch.Tensor(np.random.random((1000, 20, word_embedding_size)))  # 维度是 [ 数据量, 一句话的单词个数, 句子的embedding]
    y_train = torch.Tensor(
        np.random.randint(0, class_count, size=(1000, 1))).long()  # [ 数据量, 句子的分类], 这里的class_count=4，就是四分类任务
    return x_train, y_train


if __name__ == '__main__':
    epochs = 1000
    batch_size = 30
    embedding_size = 350
    output_class = 14
    x_train, y_train = get_total_train_data(embedding_size, output_class)
    train_loader = Data.DataLoader(
        dataset=Data.TensorDataset(x_train, y_train),  # 封装进Data.TensorDataset()类的数据，可以为任意维度
        batch_size=batch_size,  # 每块的大小
        shuffle=True,  # 要不要打乱数据 (打乱比较好)
        num_workers=6,  # 多进程（multiprocess）来读数据
        drop_last=True,
    )
    model = TextRnnModel(embedding_size=embedding_size, lstm_hidden_size=200, output_size=output_class)
    cross_loss = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)  # 优化器
    model.train()
    for i in range(epochs):
        for seq, labels in train_loader:
            optimizer.zero_grad()
            y_pred = model(seq)  # 压缩维度：得到输出，并将维度为1的去除
            single_loss = cross_loss(y_pred, labels.squeeze())
            # single_loss = cross_loss(y_pred, labels)
            single_loss.backward()
            optimizer.step()
            print("Step: " + str(i) + " loss : " + str(single_loss.detach().numpy()))