PyTorch时间序列预测GPU运行示例模型

版权声明：本文为博主原创文章，未经博主允许不得转载。 https://blog.csdn.net/Sebastien23/article/details/80574918

官方Github上给的例子是CPU版本：https://github.com/pytorch/examples/tree/master/time_sequence_prediction

下面的代码将其改为了GPU版本。除了将网络模型和输入迁移到显卡上之外，还有注意两点：

一是计算隐藏层状态和细胞状态时也需要把h和c的值迁移到显卡上；

二是在画图之前需要把预测值迁移回到CPU上才能将其转化为numpy数组。

from __future__ import print_function
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

class Sequence(nn.Module):
    def __init__(self):
        super(Sequence, self).__init__()
        self.lstm1 = nn.LSTMCell(1, 51)
        self.lstm2 = nn.LSTMCell(51, 51)
        self.linear = nn.Linear(51, 1)

    def forward(self, input, future = 0):
        outputs = []
        #h_t = torch.zeros(input.size(0), 51, dtype=torch.double)
        #c_t = torch.zeros(input.size(0), 51, dtype=torch.double)
        #h_t2 = torch.zeros(input.size(0), 51, dtype=torch.double)
        #c_t2 = torch.zeros(input.size(0), 51, dtype=torch.double)
        h_t = torch.zeros(input.size(0), 51, dtype=input.dtype, device=input.device)
        c_t = torch.zeros(input.size(0), 51, dtype=input.dtype, device=input.device)
        h_t2 = torch.zeros(input.size(0), 51, dtype=input.dtype, device=input.device)
        c_t2 = torch.zeros(input.size(0), 51, dtype=input.dtype, device=input.device)

        for i, input_t in enumerate(input.chunk(input.size(1), dim=1)):
            h_t, c_t = self.lstm1(input_t, (h_t, c_t))
            h_t2, c_t2 = self.lstm2(h_t, (h_t2, c_t2))
            output = self.linear(h_t2)
            outputs += [output]
        for i in range(future):# if we should predict the future
            h_t, c_t = self.lstm1(output, (h_t, c_t))
            h_t2, c_t2 = self.lstm2(h_t, (h_t2, c_t2))
            output = self.linear(h_t2)
            outputs += [output]
        outputs = torch.stack(outputs, 1).squeeze(2)
        return outputs


if __name__ == '__main__':
    # set random seed to 0
    np.random.seed(0)
    torch.manual_seed(0)
    # load data and make training set
    data = torch.load('traindata.pt')
    input = torch.from_numpy(data[3:, :-1])
    target = torch.from_numpy(data[3:, 1:])
    test_input = torch.from_numpy(data[:3, :-1])
    test_target = torch.from_numpy(data[:3, 1:])
    
    # move model to GPU
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print(device)
    
    input, target = input.to(device), target.to(device)
    test_input, test_target = test_input.to(device), test_target.to(device)
    
    # build the model
    seq = Sequence()
    seq.to(device)
    
    seq.double()
    criterion = nn.MSELoss()
    # use LBFGS as optimizer since we can load the whole data to train
    optimizer = optim.LBFGS(seq.parameters(), lr=0.8)
    
    epoch = 15
    #begin to train
    for i in range(epoch):
        print('STEP: ', i)
        
        def closure():
            optimizer.zero_grad()
            out = seq(input)
            #print('input size: ', input.size())
            loss = criterion(out, target)
            print('loss:', loss.item())
            loss.backward()
            return loss
    
        optimizer.step(closure)
        # begin to predict, no need to track gradient here
        with torch.no_grad():
            future = 1000
            pred = seq(test_input, future=future)
            loss = criterion(pred[:, :-future], test_target)
            print('test loss:', loss.item())
            #y = pred.detach().numpy()
            y = pred.detach().cpu().numpy()
        # draw the result
        plt.figure(figsize=(30,10))
        plt.title('Predict future values for time sequences\n(Dashlines are predicted values)', fontsize=30)
        plt.xlabel('x', fontsize=20)
        plt.ylabel('y', fontsize=20)
        plt.xticks(fontsize=20)
        plt.yticks(fontsize=20)
        def draw(yi, color):
            plt.plot(np.arange(input.size(1)), yi[:input.size(1)], color, linewidth = 2.0)
            plt.plot(np.arange(input.size(1), input.size(1) + future), yi[input.size(1):], color + ':', linewidth = 2.0)
        draw(y[0], 'r')
        draw(y[1], 'g')
        draw(y[2], 'b')
        plt.savefig('predict%d.pdf'%i)
        plt.close()