用RNN GRU逼近均值和方差

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

# 判断是否有cuda支持
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

class MeanStdRNN(nn.Module):
    def __init__(self, input_size, hidden_size):
        super(MeanStdRNN, self).__init__()
        self.rnn = nn.GRU(input_size, hidden_size, batch_first=True)
        self.linear_mean = nn.Linear(hidden_size, 1)  # predict mean
        self.linear_std = nn.Linear(hidden_size, 1)  # predict std

    def forward(self, x):
        out, _ = self.rnn(x)
        out = out[:, -1, :]  # use the output of the last step
        mean = self.linear_mean(out)
        std = self.linear_std(out)
        return mean, std

# 生成模拟数据
def generate_data(n_samples, seq_length):
    X = torch.randn(n_samples, seq_length, 1)
    y_mean = X.mean(dim=1)
    y_std = X.std(dim=1)
    return X, y_mean, y_std

n_samples = 1024
seq_length = 100
X, y_mean, y_std = generate_data(n_samples, seq_length)
X, y_mean, y_std = X.to(device), y_mean.to(device), y_std.to(device)

# 划分训练和测试数据集
split_idx = int(n_samples * 0.8)
X_train, X_test = X[:split_idx], X[split_idx:]
y_mean_train, y_mean_test = y_mean[:split_idx], y_mean[split_idx:]
y_std_train, y_std_test = y_std[:split_idx], y_std[split_idx:]

# 初始化模型、优化器和损失函数
#model = MeanStdRNN(1, 128).to(device)
#model = MeanStdRNN(1, 64).to(device)
model = MeanStdRNN(1, 32).to(device)
#optimizer = optim.Adam(model.parameters(), lr=0.001)
#optimizer = optim.SGD(model.parameters(), lr=0.01)
#optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
optimizer = torch.optim.Adamax(model.parameters(), lr=0.02)

criterion = nn.MSELoss()

# 训练模型
def train():
  rt = 1.
  n_epochs = 30
  model.train()
  for epoch in range(n_epochs):
      optimizer.zero_grad()
      mean_pred, std_pred = model(X_train)
      loss = criterion(mean_pred, y_mean_train) + criterion(std_pred, y_std_train)
      loss.backward()
      optimizer.step()
      print(f'Epoch {epoch+1}/{n_epochs}, Loss: {loss.item()}')
      rt = float(loss.item())
  return rt

from time import sleep
while train()>2**(-16): sleep(0.1)

#torch.save(model, 'meanstd_2__17.pth')

# 测试模型
def test():
  model.eval()
  with torch.no_grad():
      mean_pred, std_pred = model(X_test)
      loss = criterion(mean_pred, y_mean_test) + criterion(std_pred, y_std_test)
      print(mean_pred[0:3])
      print(y_mean_test[0:3])
      print(f'Test Loss: {loss.item()}')

test()

def draw():
  global y_mean_test,y_std_test
  model.eval()
  with torch.no_grad():
      mean_pred, std_pred = model(X_test)
      loss = criterion(mean_pred, y_mean_test) + criterion(std_pred, y_std_test)
      print(mean_pred[0:3])
      print(y_mean_test[0:3])
      print(f'Test Loss: {loss.item()}')
  import matplotlib.pyplot as plt
  # 将预测结果和实际结果从GPU移动到CPU，并转化为numpy数组
  mean_pred = mean_pred.cpu().numpy()
  std_pred = std_pred.cpu().numpy()
  y_mean_test = y_mean_test.cpu().numpy()
  y_std_test = y_std_test.cpu().numpy()
  # 画出预测的均值和实际的均值
  plt.figure(figsize=(12, 6))
  plt.subplot(1, 2, 1)
  plt.plot(mean_pred[:100], 'r', label='Predicted Mean')
  plt.plot(y_mean_test[:100], 'b', label='Actual Mean')
  plt.legend()
  # 画出预测的标准差和实际的标准差
  plt.subplot(1, 2, 2)
  plt.plot(std_pred[:100], 'r', label='Predicted Std')
  plt.plot(y_std_test[:100], 'b', label='Actual Std')
  plt.legend()
  plt.show()

draw()

when hidden 2

>>> model.state_dict()
OrderedDict([('rnn.weight_ih_l0', tensor([[-0.5121],
        [ 0.5471],
        [-0.2982],
        [ 0.2985],
        [-0.1672],
        [ 0.3252]], device='cuda:0')), ('rnn.weight_hh_l0', tensor([[ 4.4661,  0.5130],
        [ 0.0357, -0.5015],
        [-0.3948, -0.3385],
        [-0.3511,  0.3529],
        [-0.3223, -0.3423],
        [-0.0767,  1.5710]], device='cuda:0')), ('rnn.bias_ih_l0', tensor([ 0.0966, -0.3580,  1.7659,  2.0290, -0.4620, -0.0953], device='cuda:0')), ('rnn.bias_hh_l0', tensor([-0.2625,  0.8105,  2.4843,  1.7929,  1.0320,  0.0636], device='cuda:0')), ('linear_mean.weight', tensor([[-1.2655,  0.8840]], device='cuda:0')), ('linear_mean.bias', tensor([0.1250], device='cuda:0')), ('linear_std.weight', tensor([[-3.2310, -1.9122]], device='cuda:0')), ('linear_std.bias', tensor([0.2348], device='cuda:0'))])
```python