- Choose to use CPU or GPU for training
if torch.cuda.is_available():
device = torch.device("cuda")
else:
device = torch.device("cpu")
- Build LSTM model
class LSTM(nn.Module):
def __init__(self, input_size, hidden_size, num_stacked_layers,output_size):
super().__init__() # Initialize the constructor of the parent class
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) # Construct the LSTM model
self.fc = nn.Linear(hidden_size,output_size) # Fully connected layer
# Forward propagation
def forward(self, x):
batch_size = x.size(0)
# Initialize the hidden state
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)) # Separate the hidden state to avoid gradient explosion
out = self.fc(out[:, -1, :]) # Only take the last hidden state
return out
# Initialize the LSTM model
input_size=1 # Input dimension, close
hidden_size=4 # Hidden layer dimension
num_stacked_layers=1 # Number of LSTM layers
output_size=1 # Output dimension, close
model = LSTM(input_size,hidden_size,num_stacked_layers,output_size)
# model.to(device)
Parameter settings
# Define the learning rate
learning_rate = 0.001
# Define the loss function
loss_function = nn.MSELoss() # nn.CrossEntropyLoss is commonly used for binary classification problems, nn.NLLLoss is commonly used for image recognition
# Define the optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) # Choose the optimization algorithm, see https://blog.csdn.net/S20144144/article/details/103417502 for more information
- Training
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) # Forward propagation
loss = loss_function(output, y_batch) # Calculate the loss
running_loss += loss.item()
optimizer.zero_grad() # Gradient accumulation, clear the gradient
loss.backward() # Backward propagation
optimizer.step() # Update parameters
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()
- Validation
def validate_one_epoch():
model.train(False)
running_loss = 0.0
# Iterate through the test set, get data, and make predictions
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()
# Calculate
avg_loss_across_batches = running_loss / len(test_loader)
# Print
print('Val Loss: {0:.3f}'.format(avg_loss_across_batches))
- Define the number of training epochs
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()
5. Visualization
plt.plot(y_train, label='Actual Close')
plt.plot(predicted, label='Predicted Close')
plt.xlabel('Day')
plt.ylabel('Close')
plt.legend()
plt.show()