pytorch 自己写LSTM 带batch

在 PyTorch 中，可以使用 `nn.LSTM` 模块来实现 LSTM 网络。`nn.LSTM` 模块的输入数据应该是三维张量，分别表示序列长度、batch 大小和特征维度。因此，我们需要将输入数据转换为这种格式。

下面是一个示例代码，演示如何自己编写 LSTM 模型并对 batch 数据进行训练：

import torch
import torch.nn as nn

# 定义 LSTM 模型
class LSTM(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, batch_size):
        super().__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.batch_size = batch_size
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_size, 1)

    def forward(self, x):
        h0 = torch.zeros(self.num_layers, self.batch_size, self.hidden_size).to(x.device)
        c0 = torch.zeros(self.num_layers, self.batch_size, self.hidden_size).to(x.device)
        out, _ = self.lstm(x, (h0, c0))
        out = self.fc(out[:, -1, :])
        return out

# 定义超参数和数据
input_size = 10
hidden_size = 20
num_layers = 2
batch_size = 16
seq_len = 5

x = torch.randn(batch_size, seq_len, input_size)

# 将输入数据转换为三维张量
x = x.transpose(0, 1)  # 将序列长度和 batch 大小交换
x = x.transpose(0, 2)  # 将特征维度放到最前面
x = x.contiguous()     # 使张量连续

# 创建 LSTM 模型和优化器
model = LSTM(input_size, hidden_size, num_layers, batch_size)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

# 训练模型
for epoch in range(100):
    y = torch.randn(batch_size, 1)
    y_pred = model(x)
    loss = nn.functional.mse_loss(y_pred, y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    print('Epoch [{}/{}], Loss: {:.4f}'.format(epoch+1, 100, loss.item()))

在上面的代码中，我们首先定义了一个 `LSTM` 类，它包含了一个 `nn.LSTM` 模块和一个全连接层。在 `forward` 方法中，我们首先将输入数据传递给 `nn.LSTM` 模块，然后取出序列中最后一个时间步的输出，再通过全连接层得到模型的预测值。

然后，我们定义了一些超参数和随机生成的输入数据。为了将输入数据转换为三维张量，我们首先将序列长度和 batch 大小交换，然后将特征维度放到最前面，最后使用 `contiguous` 方法使张量连续。

接下来，我们创建了 `LSTM` 模型和 Adam 优化器，并进行了 100 个 epoch 的训练。在每个 epoch 中，我们随机生成了一个目标值 `y`，然后计算模型的预测值和目标值之间的均方误差，并进行反向传播和优化。