基于pytorch搭建NARX网络

以下是一个使用PyTorch搭建NARX网络的示例：

import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics import mean_squared_error

# 加载数据
data = np.load('sunspot.npy')
train_ratio = 0.7
train_size = int(train_ratio * len(data))
train_data = data[:train_size]
test_data = data[train_size:]

# 定义NARX模型
class NARX(nn.Module):
    def __init__(self):
        super(NARX, self).__init__()
        self.lstm = nn.LSTM(1, 100, batch_first=True)
        self.linear = nn.Linear(100, 1)

    def forward(self, x):
        x, _ = self.lstm(x)
        x = self.linear(x[:, -1, :])
        return x

model = NARX()
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# 训练模型
X_train, y_train = [], []
for i in range(2, len(train_data)):
    X_train.append([[train_data[i-2]], [train_data[i-1]]])
    y_train.append(train_data[i])
X_train, y_train = torch.tensor(X_train).float(), torch.tensor(y_train).float()
for epoch in range(50):
    optimizer.zero_grad()
    y_pred = model(X_train)
    loss = criterion(y_pred, y_train)
    loss.backward()
    optimizer.step()

# 测试模型
X_test, y_test = [], []
for i in range(2, len(test_data)):
    X_test.append([[test_data[i-2]], [test_data[i-1]]])
    y_test.append(test_data[i])
X_test, y_test = torch.tensor(X_test).float(), torch.tensor(y_test).float()
y_pred = model(X_test).detach().numpy()

# 绘制预测结果
plt.plot(y_test, label='Actual')
plt.plot(y_pred, label='Predicted')
plt.legend()
plt.xlabel('Month')
plt.ylabel('Sunspot Number')
plt.title('Sunspot Number Prediction using NARX')
plt.show()

# 计算均方误差
mse = mean_squared_error(y_test, y_pred)
print('Mean Squared Error:', mse)

在这个示例中，我们使用PyTorch搭建了一个NARX模型。我们使用LSTM层来处理序列数据，并使用线性层来实现NARX结构。我们使用太阳黑子数据集来训练和测试模型，并绘制了预测结果。最后，我们计算了测试集上的均方误差。