基于pytorch使用LSTM進行虎年春聯生成

先看看我摘錄的一些結果吧

七字春聯，開頭兩個字分別為 虎年、虎氣、春節

虎年嘯虎春虎虎
虎氣伏虎牛龍龍

虎年虎虎展嘯風
春節嘯春虎風虎

虎年虎嘯嘯千嘯
春節蘿嘯氣風春

虎年嘯浩一訊歡
春節回舊鶴繡舞

虎年嘯一著有處
春節回一福舞福

虎年嘯月翼業來
春節回舊精繡福

我語文水平不太行，乍一眼看感徑訓是滿高級的，

訓練資料 春聯.csv

0
龍騰虎嘯
臘盡春回
高崖伏虎嘯
茅廬臥龍飛
虎嘯風聲遠
龍騰海浪高
龍引千江水
虎越萬重山
密林藏伏虎
蘿峰染晴云
山月流古雪
風虎浴清泉
事事都如意
虎虎有生氣
云中熊虎將
天上鳳凰兒
百尺飛泉鳴震谷
一聲長嘯勢驚天
丙部琳瑯春馥郁
寅賓璀燦日光華
長白虎嘯林中日
江南猿啼嶺上風
赤縣奔騰如虎躍
神州崛起似龍飛
丑舊寅新宏圖展
牛歸虎躍春意濃
丑去寅來千里錦
牛奔虎嘯九州春
丑去寅來人益健
牛奔虎躍春愈新
春風浩蕩神州綠
虎氣升騰岳麓雄
春風著意隨人愿
虎氣生威壯國魂
春光春色源春意
虎將虎年揚虎威
春節乍聞春有喜
虎年樂見虎生風
春雷巨響山河動
月夜旋風草木飛
春曉寅回人起舞
歲於訓嘯物昭蘇
電閃金光夸五色
雷鳴巨吼動千山
憨厚忠誠牛品德
高昂奮勇虎精神
虎步奔騰開勝景
春風浩蕩展鴻圖
虎踞龍盤今勝昔
花得鳥語舊更新
虎踞龍盤今勝昔
鶯歌燕舞呈吉祥
虎年贏得春風意
喜訊喚來燕子情
虎氣頓生年屬虎
春風常駐戶迎春
虎氣頻催翻舊景
春風浩蕩著新篇
虎添雙翼前程遠
國展宏圖事業新
虎嘯大山山獻寶
龍騰祖國國揚威
虎嘯密林鳳萬壑
鶴眠蒼松月千巖
虎嘯青山千里錦
風拂綠柳萬家春
虎嘯一聲山海動
龍騰三界吉祥來
虎躍龍騰生紫氣
風調雨順兆豐年
虎躍龍騰興駿業
鶯歌燕舞羨鵬程
虎躍神州千業旺
春臨盛世萬民歡
花事才逢花好日
虎年更有虎威風
黃牛雖去精神在
猛虎初來氣象新
江山秀麗春增色
事業輝煌虎更威
江山一統騰龍日
歲月三春入虎年
皆稱飛虎一身膽
不負英雄千古名
金牛昂首高歌去
玉虎迎春斂福來
金牛辭歲寒風盡
白虎迎春喜氣來
金牛辭歲千倉滿
玉虎迎春百業興
金牛奮蹄奔大道
乳虎添翼舞新春
金牛奮蹄開錦繡
乳虎添翼會風云
金牛送舊千家樂
玉虎迎新萬戶歡
龍騰虎躍人間景
鳥語花香天地春
綠野春深禾涌碧
神州虎嘯青山來
門庭虎踞平安歲
柳浪鶯歌錦繡春
門浴春風梅吐艷
戶生虎氣鳥爭鳴
年逢寅虎群情奮
歲別丑牛大地春
牛肥馬壯豐收歲
虎躍龍騰大有年
牛肥馬壯家家富
虎躍龍騰處處春
牛奮千程榮盛世
虎馱五福賀新春
牛奮四蹄開錦繡
虎添雙翼會風云
牛耕綠野千倉滿
虎嘯青山萬木榮
牛耕沃野揚長去
虎嘯群山大步來
千載難逢新世紀
萬民譜寫虎春秋
乾元啟運三陽泰
斗丙回寅萬戶春
人逢盛世精神壯
虎躍奇峰氣勢雄
人間喜慶康平世
虎歲承歡幸福春
人民氣魄如龍虎
祖國江山似畫圖
人入虎年鼓虎勁
門添春色發春輝
人添志氣虎添翼
雪舞豐年燕舞春
人效黃牛心自貴
歲朝寅虎勁更高
山明水秀風光麗
虎躍龍騰日月新
生氣聯吟欣虎虎
留春伴讀奮年年
四海龍騰抒壯志
千山虎嘯振雄風
四海三江春氣息
千家萬戶虎精神
四海笙歌迎虎歲
九州英杰躍鵬程
唯大英雄能伏虎
是真俊杰敢擒龍
嘯一聲驚天動地
睜雙眼照耀乾坤
新年捷報虎添翼
大路朝陽馬奮蹄
興偉業仍須牛勁
展宏圖更壯虎威
一代英豪生虎氣
三春楊柳動鶯歌
英雄氣概如龍虎
祖國江山似畫圖
英雄時代英雄業
龍虎精神龍虎年
鶯歌燕舞新春日
虎躍龍騰大治年
迎春節鶯歌遍地
興中華虎勁沖天
云噴筆花騰虎豹
雨翻墨浪走蛟龍
宅后青山金虎踞
門前綠水玉龍盤
致富脫貧添虎翼
開山治水展鵬程
丙穴魚生人間改歲
寅方斗指天下皆春
春風浩蕩花香鳥語
歲月崢嶸虎躍龍騰
虎躍龍騰九州煥彩
風調雨順五谷豐登
牛奔福地普天獻瑞
虎臥華堂滿院生輝
勢如破竹人歡馬叫
安若泰山虎踞龍盤
紫氣東來江山如畫
紅旗招展龍虎揚威
祖國富強神龍活虎
人民幸福舞燕飛鶯
白虎替青牛招財進寶
黃鶯鳴翠柳辭舊迎新
虎躍龍騰創人間奇跡
鶯歌燕舞描大地春光
虎躍龍騰有天皆麗日
花香鳥語無地不春風
花團錦簇江山添異彩
虎嘯龍吟華夏壯神威
金牛辭舊攜凱歌而去
乳虎迎春帶捷報新來
瑞雪兆豐年年年大吉
丑牛接寅虎虎虎生威
歲月逢春山河添錦繡
人民思治龍虎振精神
效虎豪吟放懷歌富歲
聞雞起舞揮筆頌春光
祖國騰飛大鵬振羽翼
宏圖再展乳虎顯神通
慶虎歲把酒高吟虎躍曲
祝豐年扶犁又唱豐收謠
迎虎年敢逐改革攔路虎
送牛歲勇當奉獻老黃牛
迎新春處處呈文明氣象
入虎歲人人當改革先鋒
虎躍龍騰碧海黃山妝玉宇
鶯歌燕舞春風旭日蔚神州
虎躍龍騰華夏人民多俊杰
鶯歌燕舞陽春山水盡朝暉
牛耕廣野丑年犁出文明路
虎躍深山寅歲圖開舜堯天
歲步寅年喜慶團圓同把酒
珠還合浦歡歌一統共迎春
喜慶牛年兩制先迎香港還
歡歌兔歲億民再贊澳門歸
栽竹栽松竹隱鳳凰松隱鶴
培山培水山藏虎豹水藏龍
丑歲建奇功香港回歸昌國運
寅年興大業宏圖展現壯情懷
虎年喜虎勁攻關奪隘皆如虎
春節煥春光繡水描山總是春
牛年雖過去牛勁更增多奉獻
虎歲喜臨門虎威大振有精神
寅時入虎年十億人民振虎勁
佳節描春色九州大地蕩春潮
憶舊歲牛勁沖霄漢神鞭一指神州巨變
看今朝虎威壯中華眾志成城經濟騰飛

滑動視窗設定為2，即兩個字符預測下一個字符，滑動預測5個字符，
但是還是存在較大問題，因為我是直接將所有的對聯都拼接起來然后滑動選取訓練資料和標簽，而不是每句滑動選取，這樣就導致會串，但是懶得切了，就這樣吧，

完整代碼

# coding: utf-8
import torch
import torch.nn as nn
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.preprocessing import OneHotEncoder
import torch.nn.functional as F


class lstm_model(nn.Module):
    def __init__(self, vocab, hidden_size, num_layers, dropout=0.5):
        super(lstm_model, self).__init__()
        self.vocab = vocab  # 字符資料集
        # 索引，字符
        self.int_char = {i: char for i, char in enumerate(vocab)}
        self.char_int = {char: i for i, char in self.int_char.items()}
        # 對字符進行one-hot encoding
        self.encoder = OneHotEncoder(sparse=True).fit(vocab.reshape(-1, 1))

        self.hidden_size = hidden_size
        self.num_layers = num_layers

        # lstm層
        self.lstm = nn.LSTM(len(vocab), hidden_size, num_layers, batch_first=True, dropout=dropout)

        # 全連接層
        self.linear = nn.Linear(hidden_size, len(vocab))

    def forward(self, sequence, hs=None):
        out, hs = self.lstm(sequence, hs)  # lstm的輸出格式（batch_size, sequence_length, hidden_size）
        output = self.linear(out[:, -1])  # linear的輸出格式，(batch_size * sequence_length, vocab_size)
        return output, hs

    def onehot_encode(self, data):
        return self.encoder.transform(data)

    def onehot_decode(self, data):
        return self.encoder.inverse_transform(data)

    def label_encode(self, data):
        return np.array([self.char_int[ch] for ch in data])

    def label_decode(self, data):
        return np.array([self.int_char[ch] for ch in data])


def get_batches(data, batch_size, seq_len):
    '''
    :param data: 源資料，輸入格式(num_samples, num_features)
    :param batch_size: batch的大小
    :param seq_len: 序列的長度（精度）
    :return: （batch_size, seq_len, num_features）
    '''
    num_features = data.shape[1]
    num_chars = batch_size * seq_len  # 一個batch_size的長度

    num_batches = int(np.floor(data.shape[0] / num_chars))  # 計算出有多少個batches

    need_chars = num_batches * num_chars  # 計算出需要的總字符量

    targets = np.vstack((data[1:].A, data[0].A))  # 可能版本問題，取成numpy比較好reshape

    inputs = data[:need_chars].A.astype("int")  # 從原始資料data中截取所需的字符數量need_words

    train_data = np.zeros((inputs.shape[0] - seq_len, seq_len, num_features))
    train_label = np.zeros((inputs.shape[0] - seq_len, num_features))

    for i in range(0, inputs.shape[0] - seq_len, 1):
        train_data[i] = inputs[i:i+seq_len]
        train_label[i] = inputs[i+seq_len-1]

    for i in range(0, inputs.shape[0] - seq_len,  batch_size):
        if i + batch_size > inputs.shape[0] - seq_len:
            break
        x = train_data[i:i+batch_size]
        y = train_label[i:i+batch_size]
        yield x, y


def train(model, data, batch_size, seq_len, epochs, lr=0.01, valid=None):
    device = 'cuda' if torch.cuda.is_available() else 'cpu'

    model = model.to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()

    if valid is not None:
        data = model.onehot_encode(data.reshape(-1, 1))
        valid = model.onehot_encode(valid.reshape(-1, 1))
    else:
        data = model.onehot_encode(data.reshape(-1, 1))

    train_loss = []
    val_loss = []

    for epoch in range(epochs):
        model.train()
        hs = None  # hs等于hidden_size隱藏層節點
        train_ls = 0.0
        val_ls = 0.0
        for x, y in get_batches(data, batch_size, seq_len):
            optimizer.zero_grad()
            x = torch.tensor(x).float().to(device)
            out, hs = model(x, hs)
            hs = ([h.data for h in hs])
            y = y.reshape(-1, len(model.vocab))
            y = model.onehot_decode(y)
            y = model.label_encode(y.squeeze())
            y = torch.from_numpy(y).long().to(device)

            loss = criterion(out, y.squeeze())
            loss.backward()
            optimizer.step()
            train_ls += loss.item()

        if valid is not None:
            model.eval()
            hs = None
            with torch.no_grad():
                for x, y in get_batches(valid, batch_size, seq_len):
                    x = torch.tensor(x).float().to(device)  # x為一組測驗資料，包含batch_size * seq_len個字
                    out, hs = model(x, hs)

                    # out.shape輸出為tensor[batch_size * seq_len, vocab_size]
                    hs = ([h.data for h in hs])  # 更新引數

                    y = y.reshape(-1, len(model.vocab))  # y.shape為(128,100,43)，因此需要轉成兩維，每行就代表一個字了，43為字典大小
                    y = model.onehot_decode(y)  # y標簽即為測驗資料各個字的下一個字，進行one_hot解碼，即變為字符
                    # 但是此時y 是[[..],[..]]形式
                    y = model.label_encode(y.squeeze())  # 因此需要去掉一維才能成功解碼
                    # 此時y為[12...]成為一維的陣列，每個代表自己字典里對應字符的字典序
                    y = torch.from_numpy(y).long().to(device)

                    # 這里y和y.squeeze()出來的東西一樣，可能這里沒啥用，不太懂
                    loss = criterion(out, y.squeeze())  # 計算損失值
                    val_ls += loss.item()

            val_loss.append(np.mean(val_ls))
        train_loss.append(np.mean(train_ls))
        print("val_loss", val_ls)
        print("train_loss:", train_ls)

    plt.plot(train_loss, label="train_loss")
    plt.plot(val_loss, label="val loss")
    plt.title("loop vs epoch")
    plt.legend()
    plt.show()

    model_name = "lstm_model.net"

    with open(model_name, 'wb') as f:  # 訓練完了保存模型
        torch.save(model.state_dict(), f)


def predict(vocab_size, model, char, top_k=None, hidden_size=None):
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    model.to(device)
    model.eval()  # 固定引數
    with torch.no_grad():
        char = np.array([char])  # 輸入一個字符，預測下一個字是什么，先轉成numpy
        char = char.reshape(-1, 1)  # 變成二維才符合編碼規范
        char_encoding = model.onehot_encode(char).A  # 對char進行編碼，取成numpy比較方便reshape
        char_encoding = char_encoding.reshape(1, -1, vocab_size)# 轉成模型輸入格式 

        char_tensor = torch.tensor(char_encoding, dtype=torch.float32)  # 轉成tensor
        char_tensor = char_tensor.to(device)

        out, hidden_size = model(char_tensor, hidden_size)  # 放入模型進行預測，out為結果

        probs = F.softmax(out, dim=1).squeeze()  # 計算預測值,即所有字符的概率
        print(probs.shape)

        if top_k is None:  # 選擇概率最大的top_k個
            indices = np.arange(vocab_size)
        else:
            probs, indices = probs.topk(top_k)
            indices = indices.cpu().numpy()
        probs = probs.cpu().numpy()

        char_index = np.random.choice(indices, p=probs/probs.sum())  # 隨機選擇一個字符索引作為預測值
        char = model.int_char[char_index]  # 通過索引找出預測字符

    return char, hidden_size


def sample(vocab_size, seq_len, model, length,sentence, top_k=None):
    hidden_size = None
    new_sentence = [char for char in sentence]
    for i in range(length):

        next_char, hidden_size = predict(vocab_size, model, new_sentence[-seq_len:], top_k=top_k, hidden_size=hidden_size)

        new_sentence.append(next_char)
    return "".join(new_sentence)


def main():
    hidden_size = 512
    num_layers = 4
    batch_size = 128
    seq_len = 2
    epochs = 30
    lr = 0.0001

    f = pd.read_csv("春聯.csv")
    f = f["0"]
    text = list(f)
    text = ".".join(text).replace(".", "")
    vocab = np.array(sorted(set(text)))  # 建立字典
    vocab_size = len(vocab)
    print("vocab_size", vocab_size)


    val_len = int(np.floor(0.2 * len(text)))  # 劃分訓練測驗集
    trainset = np.array(list(text[:-val_len]))
    validset = np.array(list(text[-val_len:]))

    model = lstm_model(vocab, hidden_size, num_layers)  # 模型實體化
    train(model, trainset, batch_size, seq_len, epochs, lr=lr, valid=validset)  # 訓練模型
    model.load_state_dict(torch.load("lstm_model.net"))  # 呼叫保存的模型
    new_text1 = sample(vocab_size, seq_len, model, 5, "虎年", top_k=7)  # 預測模型，生成100個字符,預測時選擇概率最大的前5個
    new_text2 = sample(vocab_size, seq_len, model, 5, "春節", top_k=7)  # 預測模型，生成100個字符,預測時選擇概率最大的前5個
    print(new_text1)  # 輸出預測文本
    print(new_text2)


if __name__ == "__main__":
    main()