基于Pytorch的LSTM网络全流程实验（自带数据集，可直接运行出结果，替换自己的数据集即可使用）

import numpy as np
import matplotlib.pyplot as plt

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset

from sklearn.model_selection import train_test_split


# 生成数据集
def generate_data(num_samples, seq_length):
    # 生成正弦波形数据（类别0）
    half_num_samples=num_samples//2 # 整除
    x_sin = np.array([np.sin(0.06 * np.arange(seq_length) + np.random.rand()) for _ in range(half_num_samples)])
    y_sin = np.zeros(half_num_samples, dtype=np.int64)
    
    # 生成余弦波形数据（类别1），频率略有不同
    x_cos = np.array([np.cos(0.05 * np.arange(seq_length) + np.random.rand()) for _ in range(half_num_samples)])
    y_cos = np.ones(half_num_samples, dtype=np.int64)
    
    # 合并数据
    x = np.concatenate((x_sin, x_cos), axis=0)
    y = np.concatenate((y_sin, y_cos), axis=0)
    
    # 打乱数据
    indices = np.arange(num_samples)
    np.random.shuffle(indices)
    x = x[indices]
    y = y[indices]
    
    # 转换为pytorch张量，LSTM需要3D tensor [batch, seq_len, features]，
    # 所以用unsqueeze(2)在第二个维度上增加一个维度
    x_tensor = torch.tensor(x, dtype=torch.float32).unsqueeze(2)  
    print('x_tensor.shape:',x_tensor.shape) # x_tensor.shape: torch.Size([1000, 100, 1])
    y_tensor = torch.tensor(y, dtype=torch.int64) # y_tensor.shape: torch.Size([1000])
    print('y_tensor.shape:',y_tensor.shape)
    
    return x_tensor, y_tensor

# LSTM分类模型
class LSTMClassifier(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim, n_layers):
        super(LSTMClassifier, self).__init__()
        self.hidden_dim = hidden_dim
        self.n_layers = n_layers
        
        # LSTM Layer
        self.lstm = nn.LSTM(input_dim, hidden_dim, n_layers, batch_first=True)
        
        # 全连接层（Fully connected layer）
        self.fc = nn.Linear(hidden_dim, output_dim)
    
    # forward方法在模型训练时会自动调用
    def forward(self, x):
        # 用零初始化隐藏层的状态
        h0 = torch.zeros(self.n_layers, x.size(0), self.hidden_dim).requires_grad_()
        
        # 用零初始化细胞状态
        c0 = torch.zeros(self.n_layers, x.size(0), self.hidden_dim).requires_grad_()
        
        out, (hn, cn) = self.lstm(x, (h0.detach(), c0.detach()))
        
        out = self.fc(out[:, -1, :])
        return out

# 训练模型
def train_model(model, train_loader, criterion, optimizer, num_epochs):
    for epoch in range(num_epochs):
        for i, (sequences, labels) in enumerate(train_loader):
            # Forward pass
            outputs = model(sequences)
            loss = criterion(outputs, labels)
            
            # Backward and optimize
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
            if (i+1) % 100 == 0:
                print(f'Epoch [{epoch+1}/{num_epochs}], Step [{i+1}], Loss: {loss.item():.4f}')

# 评估模型
def evaluate_model(model, test_loader):
    model.eval()  # Set model to evaluation mode
    with torch.no_grad():
        correct = 0
        total = 0
        for sequences, labels in test_loader:
            outputs = model(sequences)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
            
        print(f'Accuracy of the model on the test sequences: {100 * correct / total} %')




if __name__=='__main__':
    # ----------------- 生成样本数据 ----------------- 
    num_samples = 1000  # 训练总样本数
    seq_length = 100    # 每个样本的序列长度（可以看作是特征的长度）
    x_data,y_data = generate_data(num_samples, seq_length) # 产生总的样本
    x_train, x_test, y_train, y_test = train_test_split(x_data, y_data, test_size=0.2, random_state=2) 
    # ----------------- 数据加载器 ----------------- 
    batch_size=64
    train_loader = DataLoader(TensorDataset(x_train, y_train), batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(TensorDataset(x_test, y_test), batch_size=batch_size, shuffle=False)

    # ----------------- 可视化数据 ----------------- 
    plt.figure(figsize=(12, 6))
    for i in range(6):
        plt.subplot(2, 3, i+1)
        plt.plot(x_train[i].numpy().flatten(), label=f"Class {y_train[i].item()}")
        plt.legend()
    plt.tight_layout()
    plt.show() # 不想看数据，可以注释掉这行

    # ----------------- 超参设定 ----------------- 
    input_dim = 10    # 输入特征的维数
    hidden_dim = 50   # LSTM 隐藏层的维度
    output_dim = 2    # 输出的维度（分类的类别数）
    n_layers = 1      # 堆叠的 LSTM 层的数量（默认为1层）

    # ----------------- 创建模型 ----------------- 
    model = LSTMClassifier(input_dim=1, hidden_dim=50, output_dim=2, n_layers=1)
    criterion = nn.CrossEntropyLoss() # 损失函数
    optimizer = optim.Adam(model.parameters(), lr=0.01) # 优化器

    # ----------------- 训练模型 ----------------- 
    train_model(model, train_loader, criterion, optimizer, num_epochs=10)

    # ----------------- 评估模型 ----------------- 
    evaluate_model(model, test_loader)