用pytorch写一个lstm进行轨迹预测的代码

好的，以下是一个简单的基于PyTorch的LSTM轨迹预测代码示例：

import torch
import torch.nn as nn
import numpy as np

# 定义LSTM模型
class LSTM(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(LSTM, self).__init__()
        self.hidden_size = hidden_size
        self.lstm = nn.LSTM(input_size, hidden_size)
        self.fc = nn.Linear(hidden_size, output_size)
        
    def forward(self, input):
        lstm_out, _ = self.lstm(input.view(len(input), 1, -1))
        output = self.fc(lstm_out[-1])
        return output

# 随机生成轨迹数据
data = np.random.randn(100, 5)

# 将数据分为训练集和测试集
train_data = data[:80]
test_data = data[80:]

# 定义模型和优化器
input_size = 5
hidden_size = 10
output_size = 5
model = LSTM(input_size, hidden_size, output_size)
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

# 训练模型
for epoch in range(100):
    train_loss = 0
    for i in range(len(train_data) - 1):
        optimizer.zero_grad()
        input = torch.Tensor(train_data[i])
        target = torch.Tensor(train_data[i+1])
        output = model(input)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        train_loss += loss.item()
    print('Epoch: {}, Loss: {:.6f}'.format(epoch+1, train_loss/len(train_data)))

# 测试模型
test_loss = 0
with torch.no_grad():
    for i in range(len(test_data) - 1):
        input = torch.Tensor(test_data[i])
        target = torch.Tensor(test_data[i+1])
        output = model(input)
        loss = criterion(output, target)
        test_loss += loss.item()
print('Test Loss: {:.6f}'.format(test_loss/len(test_data)))

这个例子中，我们使用了一个包含5个维度的随机轨迹数据。我们将这些数据分为训练集和测试集，然后定义了一个LSTM模型，并使用均方误差损失函数和Adam优化器进行训练。在训练过程中，我们每个时刻输入一个时间步的数据，并预测下一个时间步的数据。在测试过程中，我们使用训练好的模型来预测测试集中的数据，并计算模型的平均测试损失。