pyorch 通過前三個數預測後一個數的模型代碼demo

相關知識：

• 使用scatter_進行one-hot編碼

scatter_(self, dim, index, value)將src中數據根據index中的索引按照dim的方向填入value值，這個函數可以從轉換成onehot編碼來理解

• LSTM返回output、hidden和cell

output, (hidden,cell) = self.LSTM(x)

• RNN返回output、hidden

output, hidden = self.RNN(x)

 

# -----------------------------------
# 模塊導入
import numpy
import torch
from torch import nn

# -----------------------------------
# 數據預處理
data_length = 30
# 定義30個數，通過前三個預測後一個，比如：1,2,3->4
seq_length = 3
# 通過上面可知序列長度爲3

number = [i for i in range(data_length)]
li_x = []
li_y = []
for i in range(0, data_length - seq_length):
    x = number[i: i + seq_length]
    y = number[i + seq_length]
    li_x.append(x)
    li_y.append(y)
#number: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29]
#li_x: [[0, 1, 2], [1, 2, 3], [2, 3, 4], [3, 4, 5], [4, 5, 6], [5, 6, 7], [6, 7, 8], [7, 8, 9], [8, 9, 10],
#       [9, 10, 11], [10, 11, 12], [11, 12, 13], [12, 13, 14], [13, 14, 15], [14, 15, 16], [15, 16, 17], [16, 17, 18], [17, 18, 19],
#       [18, 19, 20], [19, 20, 21], [20, 21, 22], [21, 22, 23], [22, 23, 24], [23, 24, 25], [24, 25, 26], [25, 26, 27], [26, 27, 28]]
#li_y: [3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29]

# 輸入數據格式：seq_len, batch, input_size (27,1,3)
data_x = numpy.reshape(li_x, (len(li_x), 1, seq_length))

# 將輸入數據歸一化
data_x = torch.from_numpy(data_x / float(data_length)).float()

# scatter_函數使用來轉換onehot編碼
# 將輸出數據設置爲one-hot編碼
data_y = torch.zeros(len(li_y), data_length).scatter_(1, torch.tensor(li_y).unsqueeze_(dim=1), 1).float()


# -----------------------------------
# 定義網絡模型
class net(nn.Module):
    # 模型結構：LSTM + 全連接 + Softmax
    def __init__(self, input_size, hidden_size, output_size, num_layer):
        super(net, self).__init__()
        # LSTM返回output、hidden和cell
        # RNN返回output、hidden
        self.layer1 = nn.LSTM(input_size, hidden_size, num_layer)
        self.layer2 = nn.Linear(hidden_size, output_size)
        self.layer3 = nn.Softmax()

    def forward(self, x):
        x, _ = self.layer1(x)
        # 格式：[27, 1, 32]，代表樣本數量，batch大小以及隱藏層尺寸
        sample, batch, hidden = x.size()
        x = x.reshape(-1, hidden)
        # 轉成二維矩陣後與全連接進行計算
        x = self.layer2(x)
        x = self.layer3(x)
        return x


model = net(seq_length, 32, data_length, 4)

# -----------------------------------
# 定義損失函數和優化器
loss_fun = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

# -----------------------------------
# 訓練模型

# 訓練前可以先看看初始化的參數預測的結果差距
# result = model(data_x)
# for target, pred in zip(data_y, result):
#     print("{} -> {}".format(target.argmax().data, pred.argmax().data))

# 開始訓練1000輪
for _ in range(1000):
    output = model(data_x)
    loss = loss_fun(data_y, output)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if (_ + 1) % 50 == 0:
        print('Epoch: {}, Loss: {}'.format(_, loss.data))

# -----------------------------------
# 預測結果
result = model(data_x)
for target, pred in zip(data_y, result):
    print("正確結果：{}，預測：{}".format(target.argmax().data, pred.argmax().data))

precision = (result.argmax(dim=1).data==data_y.argmax(dim=1).data)


    # 結果：
    # 正確結果：3，預測：3
    # 正確結果：4，預測：4
    # 正確結果：5，預測：5
    # 正確結果：6，預測：6
    # 正確結果：7，預測：7
    # 正確結果：8，預測：8
    # 正確結果：9，預測：9
    # 正確結果：10，預測：10
    # 正確結果：11，預測：11
    # 正確結果：12，預測：12
    # 正確結果：13，預測：13
    # 正確結果：14，預測：14
    # 正確結果：15，預測：15
    # 正確結果：16，預測：16
    # 正確結果：17，預測：21
    # 正確結果：18，預測：18
    # 正確結果：19，預測：27
    # 正確結果：20，預測：21
    # 正確結果：21，預測：21
    # 正確結果：22，預測：21
    # 正確結果：23，預測：21
    # 正確結果：24，預測：24
    # 正確結果：25，預測：25
    # 正確結果：26，預測：26
    # 正確結果：27，預測：27
    # 正確結果：28，預測：28
    # 正確結果：29，預測：29