Pytorch官方教程：用RNN实现字符级的分类任务

数据处理

数据可以从传送门下载。

这些数据包括了18个国家的名字，我们的任务是根据这些数据训练模型，使得模型可以判断出名字是哪个国家的。

一开始，我们需要对名字进行一些处理，因为不同国家的文字可能会有一些区别。

在这里最好先了解一下Unicode：可以看看：Unicode的文本处理二三事

                                                                       NLP哪里跑: Unicode相关的一些小知识和工具

import os
import glob
import unicodedata

all_letters = string.ascii_letters + " .,;'"  # string.ascii_letters的作用是生成所有的英文字母
n_letters = len(all_letters)

def find_files(path):
    """
    :param path:文件路径
    :return: 文件列表地址
    """
    return glob.glob(path)     # glob模块提供了一个函数用于从目录通配符搜索中生成文件列表:


def unicode_to_ascii(str):
    """
    :param str:名字
    :return:返回均采用NFD编码方式的名字
    """
    return ''.join(
      c for c in unicodedata.normalize('NFD', str)    # 对文字采用相同的编码方式
      if unicodedata.category(c) != 'Mn' and c in all_letters 
    )


def read_lines(files_list):
    """
    读取每个文件的内容
    :param files_list:文件所在地址列表
    :return:{国家：名字列表}
    """
    category_lines = {}
    all_categories = []
    for file in files_list:
        # os.path.splitext:分割路径，返回路径名和文件扩展名的元组
        # os.path.basename:返回文件名
        category = os.path.splitext(os.path.basename(file))[0]
        line = [unicode_to_ascii(line) for line in open(file)]
        all_categories.append(category)
        category_lines[category] = line
    return all_categories, category_lines
    # print(all_categories)
    # print(category_lines['Chinese'])

接下来我们要对单词进行编码，这里使用独热编码方式， 在上述代码中已经生成了all_letters的字符串，对于名字中的每个字母，我们只需令其在all_letters中所在的索引位为0即可。

这样每个名字的size就是[ len(name)，1，len(all_letters) ]。

def get_index(letter):
    """
    :param letter: 字母
    :return: 字母索引
    """
    return all_letters.find(letter)


def letter_to_tensor(letter):
    """
    将字母转换成张量
    :param letter:字母
    :return: 张量
    """
    tensor = torch.zeros(1, n_letters)
    tensor[0][get_index(letter)] = 1
    return tensor.to(device)    # 将tensor放到cuda上


def word_to_tensor(word):
    """
    将单词转换成张量
    :param word: 单词
    :return: 张量
    """
    tensor = torch.zeros(len(word), 1, n_letters)
    for i, letter in enumerate(word):
        tensor[i][0][get_index(letter)] = 1
    return tensor.to(device)

　　

模型构建

因为是入门级的学习，这里也是只使用了最简单的RNN，只包含了一个隐藏层。

这里将隐藏层的维度定为128。

class RNN(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(RNN, self).__init__()

        self.input_size = input_size
        self.hidden_size = hidden_size
        self.output_size = output_size

        self.i2h = nn.Linear(input_size + self.hidden_size, self.hidden_size)
        self.i2o = nn.Linear(input_size + self.hidden_size, self.output_size)  # 这里写成self.i2o = nn.Linear(self.hidden_size, self.output_size)也是可以的，forward传入的就得是hidden的值 
        self.softmax = nn.LogSoftmax(dim=1)  # dim=1表示对第1维度的数据进行logsoftmax操作

    def forward(self, input, hidden):
        tmp = torch.cat((input, hidden), 1)   
        hidden = self.i2h(tmp)
        output = self.i2o(tmp)
        output = self.softmax(output)
        return output, hidden

    def init_hidden(self):   # 隐藏层初始化0操作
        return torch.zeros(1, self.hidden_size).to(device)

最后输出的结果是属于18个国家的概率值，下面的函数就是从中挑选出概率最大的国家。

def category_from_output(output):
    top_n, top_i = output.topk(1)
    category_i = top_i[0].item()
    return all_categories[category_i], category_i

　　

模型训练

模型的训练采用随机梯度下降，每次随机选择一个数据来进行训练。下面的代码实现的功能是先随机选择一个国家，再从该国家中随机选择一个名字。

def random_choice(obj):
    return obj[random.randint(0, len(obj)-1)]


def random_training_example():
    category = random_choice(all_categories)
    word = random_choice(category_lines[category])
    category_tensor = torch.tensor([all_categories.index(category)], dtype=torch.long).to(device)
    word_tensor = word_to_tensor(word)
    return category, word, category_tensor, word_tensor

具体的训练代码如下：

def train():
    rnn = RNN(n_letters, n_hidden, n_categories).to(device)   # 创建模型
    loss = nn.NLLLoss()                                      # 定义损失函数
    lr = 0.005                                               # 学习率
    epoch_num = 100000                                       # 迭代次数
    current_loss = 0                                         # 累计损失
    all_losses = []                                          # 记录损失，后续画图

    for epoch in range(epoch_num):
        epoch_start_time = time.time()                       # 记录一次迭代的时间
        category, word, category_tensor, word_tensor = random_training_example()     # 随机选择一条训练数据

        rnn.zero_grad()                                      # 梯度清零，这和optimizer.zero_grad()是等价的
        hidden = rnn.init_hidden()                           # 初始化隐藏层
        for i in range(word_tensor.size()[0]):
            output, hidden = rnn(word_tensor[i].to(device), hidden.to(device))

        train_loss = loss(output, category_tensor)
        train_loss.backward()

        for p in rnn.parameters():
            p.data.add_(p.grad.data, alpha=-lr)

        current_loss += train_loss.item()

        if epoch % 5000 == 0:                                # 每迭代5000次输出信息
            guess, guess_i = category_from_output(output)
            correct = '√' if guess == category else '×(%s)' % category
            print('%d %d%% %2.4f sec(s) %.4f %s / %s %s' %
                  (epoch, epoch / epoch_num * 100, time.time() - epoch_start_time, train_loss.item(), word, guess, correct))

        if (epoch+1) % 1000 == 0:                            # 每迭代1000次，记录下该1000次的平均损失
            all_losses.append(current_loss / 1000)
            current_loss = 0
    plt.figure()                                             # 画出损失变化图
    plt.plot(all_losses)
    plt.show()

最终的损失变化图为：

完整代码

import string
import os
import glob
import unicodedata
import torch
import torch.nn as nn
import random
import time
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker


device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

all_letters = string.ascii_letters + " .,;'"  # string.ascii_letters的作用是生成所有的英文字母
n_letters = len(all_letters)
n_hidden = 128
n_categories = 18


class RNN(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(RNN, self).__init__()

        self.input_size = input_size
        self.hidden_size = hidden_size
        self.output_size = output_size

        self.i2h = nn.Linear(input_size + self.hidden_size, self.hidden_size)
        self.i2o = nn.Linear(input_size + self.hidden_size, self.output_size)    # 这里写成self.i2o = nn.Linear(self.hidden_size, self.output_size)也是可以的，forward传入的就得是hidden的值
        self.softmax = nn.LogSoftmax(dim=1)  # dim=1表示对第1维度的数据进行logsoftmax操作

    def forward(self, input, hidden):
        tmp = torch.cat((input, hidden), 1)
        hidden = self.i2h(tmp)
        output = self.i2o(tmp)
        output = self.softmax(output)
        return output, hidden

    def init_hidden(self):   # 隐藏层初始化0操作
        return torch.zeros(1, self.hidden_size).to(device)


def find_files(path):
    """
    :param path:文件路径
    :return: 文件列表地址
    """
    return glob.glob(path)     # glob模块提供了一个函数用于从目录通配符搜索中生成文件列表:


def unicode_to_ascii(str):
    """
    :param str:名字
    :return:返回均采用NFD编码方式的名字
    """
    return ''.join(
      c for c in unicodedata.normalize('NFD', str)    # 对文字采用相同的编码方式
      if unicodedata.category(c) != 'Mn' and c in all_letters 
    )


def read_lines(files_list):
    """
    读取每个文件的内容
    :param files_list:文件所在地址列表
    :return:{国家：名字列表}
    """
    category_lines = {}
    all_categories = []
    for file in files_list:
        # os.path.splitext:分割路径，返回路径名和文件扩展名的元组
        # os.path.basename:返回文件名
        category = os.path.splitext(os.path.basename(file))[0]
        line = [unicode_to_ascii(line) for line in open(file)]
        all_categories.append(category)
        category_lines[category] = line
    return all_categories, category_lines
    # print(all_categories)
    # print(category_lines['Chinese'])


def get_index(letter):
    """
    :param letter: 字母
    :return: 字母索引
    """
    return all_letters.find(letter)


def letter_to_tensor(letter):
    """
    将字母转换成张量
    :param letter:字母
    :return: 张量
    """
    tensor = torch.zeros(1, n_letters)
    tensor[0][get_index(letter)] = 1
    return tensor.to(device)    # 将tensor放到cuda上


def word_to_tensor(word):
    """
    将单词转换成张量
    :param word: 单词
    :return: 张量
    """
    tensor = torch.zeros(len(word), 1, n_letters)
    for i, letter in enumerate(word):
        tensor[i][0][get_index(letter)] = 1
    return tensor.to(device)


def random_choice(obj):
    return obj[random.randint(0, len(obj)-1)]


def random_training_example():
    category = random_choice(all_categories)
    word = random_choice(category_lines[category])
    category_tensor = torch.tensor([all_categories.index(category)], dtype=torch.long).to(device)
    word_tensor = word_to_tensor(word)
    return category, word, category_tensor, word_tensor


def category_from_output(output):
    top_n, top_i = output.topk(1)
    category_i = top_i[0].item()
    return all_categories[category_i], category_i


def train():
    rnn = RNN(n_letters, n_hidden, n_categories).to(device)   # 创建模型
    loss = nn.NLLLoss()                                      # 定义损失函数
    lr = 0.005                                               # 学习率
    epoch_num = 100000                                       # 迭代次数
    current_loss = 0                                         # 累计损失
    all_losses = []                                          # 记录损失，后续画图

    for epoch in range(epoch_num):
        epoch_start_time = time.time()                       # 记录一次迭代的时间
        category, word, category_tensor, word_tensor = random_training_example()     # 随机选择一条训练数据

        rnn.zero_grad()                                      # 梯度清零，这和optimizer.zero_grad()是等价的
        hidden = rnn.init_hidden()                           # 初始化隐藏层
        for i in range(word_tensor.size()[0]):
            output, hidden = rnn(word_tensor[i].to(device), hidden.to(device))

        train_loss = loss(output, category_tensor)
        train_loss.backward()

        for p in rnn.parameters():
            p.data.add_(p.grad.data, alpha=-lr)

        current_loss += train_loss.item()

        if epoch % 5000 == 0:                                # 每迭代5000次输出信息
            guess, guess_i = category_from_output(output)
            correct = '√' if guess == category else '×(%s)' % category
            print('%d %d%% %2.4f sec(s) %.4f %s / %s %s' %
                  (epoch, epoch / epoch_num * 100, time.time() - epoch_start_time, train_loss.item(), word, guess, correct))

        if (epoch+1) % 1000 == 0:                            # 每迭代1000次，记录下该1000次的平均损失
            all_losses.append(current_loss / 1000)
            current_loss = 0
    plt.figure()                                             # 画出损失变化图
    plt.plot(all_losses)
    plt.show()


if __name__ == '__main__':
    files_list = find_files('./names/*.txt')
    all_categories, category_lines = read_lines(files_list)
    train()