skka3134

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

class LSTM(nn.Module):
    def __init__(self, input_size, hidden_size, num_stacked_layers,output_size):
        super().__init__() #初始化父类中的构造方法
        self.hidden_size = hidden_size
        self.num_stacked_layers = num_stacked_layers
        self.lstm = nn.LSTM(input_size, hidden_size, num_stacked_layers, batch_first=True)#构造lstm模型
        self.fc = nn.Linear(hidden_size,output_size) #全连接层
    #前向传播
    def forward(self, x):
        batch_size = x.size(0)
        #初始化隐藏层状态
        h0 = torch.zeros(self.num_stacked_layers, batch_size, self.hidden_size).to(device)
        c0 = torch.zeros(self.num_stacked_layers, batch_size, self.hidden_size).to(device)
        out, _ = self.lstm(x, (h0, c0)) #分离隐藏状态，以免梯度爆炸
        out = self.fc(out[:, -1, :]) #只要最后一层隐层状态
        return out 

#初始化lstm模型
input_size=1 #输入维度，close
hidden_size=4 #隐藏层维度
num_stacked_layers=1  #lstm层数
output_size=1 #输出维度，close
model = LSTM(input_size,hidden_size,num_stacked_layers,output_size)
# model.to(device)

#定义学习率
learning_rate = 0.001
#定义损失函数
loss_function = nn.MSELoss() #nn.CrossEntropyLoss常用于解决二分类问题，nn.NLLLoss常用于图像识别
#定义优化器
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) #关于优化算法的选择https://blog.csdn.net/S20144144/article/details/103417502

def train_one_epoch():
    model.train(True)
    print(f'Epoch: {epoch + 1}')
    running_loss = 0.0
    
    for batch_index, batch in enumerate(train_loader):
        x_batch, y_batch = batch[0].to(device), batch[1].to(device)

        output = model(x_batch)#前向传播
        loss = loss_function(output, y_batch)#计算损失
        running_loss += loss.item()
        optimizer.zero_grad() #梯度会一直累加，清零梯度。
        loss.backward() #反向传播
        optimizer.step() #更新参数

        if batch_index % 100 == 99:  # print every 100 batches
            avg_loss_across_batches = running_loss / 100
            print('Batch {0}, Loss: {1:.3f}'.format(batch_index+1,avg_loss_across_batches))
            running_loss = 0.0
    print()

def validate_one_epoch():
    model.train(False)
    running_loss = 0.0
    #迭代测试集，获取数据，预测
    for batch_index, batch in enumerate(test_loader):
        x_batch, y_batch = batch[0].to(device), batch[1].to(device)
        with torch.no_grad():
            output = model(x_batch)
            loss = loss_function(output, y_batch)
            running_loss += loss.item()
    #计算
    avg_loss_across_batches = running_loss / len(test_loader)
    #打印
    print('Val Loss: {0:.3f}'.format(avg_loss_across_batches))

num_epochs = 10
for epoch in range(num_epochs):
    train_one_epoch()
    validate_one_epoch()
with torch.no_grad():
    predicted = model(X_train.to(device)).to('cpu').numpy()

plt.plot(y_train, label='Actual Close')
plt.plot(predicted, label='Predicted Close')
plt.xlabel('Day')
plt.ylabel('Close')
plt.legend()
plt.show()