神经网络预测地震加速度反应谱曲线，pytorch实现

首先要下载好的库：pytorch, matplotlib
推荐使用pycharm + Anaconda的环境配置

一、生成数据集

data2.py

import numpy as np
# 特征周期Tg
a = []
for i in range(10):
    c = i + 1
    if c <= 6:
        a.append(round(0.2 + (c - 1) * 0.05, 2))
    elif c <= 9:
        a.append(round(0.55 + (c - 7) * 0.1, 2))
    else:
        a.append(round(0.9, 2))

# 阻尼比theta
c = []
for i in range(6):
    if i == 0:
        theta = 0.01
    elif i == 1:
        theta = 0.02
    elif i == 2:
        theta = 0.05
    elif i == 3:
        theta = 0.1
    elif i == 4:
        theta = 0.2
    elif i == 5:
        theta = 0.3

    c.append(theta)


# 输入矩阵，[Tg, theta, T, alpha]
input_data = []
for i in range(10):
    for k in range(6):
        for j in range(12):
            for l in range(600):
                input_data.append([a[i], c[k], round(0.01 * l, 2)])

                if j == 0:
                    alpha = 0.04
                elif j == 1:
                    alpha = 0.08
                elif j == 2:
                    alpha = 0.12
                elif j == 3:
                    alpha = 0.16
                elif j == 4:
                    alpha = 0.24
                elif j == 5:
                    alpha = 0.28
                elif j == 6:
                    alpha = 0.32
                elif j == 7:
                    alpha = 0.5
                elif j == 8:
                    alpha = 0.72
                elif j == 9:
                    alpha = 0.9
                elif j == 10:
                    alpha = 1.2
                elif j == 11:
                    alpha = 1.4

                input_data[-1].append(alpha)

target_data = []

for i in range(10):
    for k in range(6):
        for j in range(12):
            for l in range(600):
                tem = input_data[i * 6 * 12 * 600 + k * 12 * 600 + j * 600 + l]


                Tg1 = tem[0]
                theta1 = tem[1]
                T1 = tem[2]
                alpha1 = tem[3]

                gamma = round(0.9 + (0.05 - theta1) / (0.3 + 6 * theta1), 2)
                eta1 = round(0.02 + (0.05 - theta1) / (4 + 32 * theta1), 3)
                eta2 = round(1 + (0.05 - theta1) / (0.06 + 1.6 * theta1), 2)

                if T1 <= 0.1:
                    target_data.append([round(10 * (eta2 - 0.45) * alpha1 * T1 + 0.45 * alpha1, 5)])
                elif T1 <= Tg1:
                    target_data.append([round(eta2 * alpha1, 5)])
                elif T1 <= 5 * Tg1:
                    target_data.append([round(pow((Tg1 / T1), gamma) * eta2 * alpha1, 5)])
                else:
                    target_data.append([round((eta2 * pow(0.2, gamma) - eta1 * (T1 - 5 * Tg1)) * alpha1, 5)])


input_data = np.array(input_data)
target_data = np.array(target_data)

# print(input_data)
# print(len(input_data))
#
# print(target_data)
# print(len(target_data))

二、搭建神经网络

main.py

import torch
from torch import nn
from torch.utils.data import DataLoader
import torch.utils.data as Data
import data2

# 数据集
input = torch.from_numpy(data2.input_data)
target = torch.from_numpy(data2.target_data)

input = torch.tensor(input, dtype=torch.float32)
target = torch.tensor(target, dtype=torch.float32)

# print(input)
# print(target)

# 搭建神经网络
model = torch.nn.Sequential(
    torch.nn.Linear(1*4, 500),
    torch.nn.ReLU(),
    torch.nn.Linear(500, 500),
    torch.nn.ReLU(),
    torch.nn.Linear(500, 500),
    torch.nn.ReLU(),
    torch.nn.Linear(500, 1)

)

# 损失函数和优化器
loss_fuc = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# 数据集和下载器
train_dataset = Data.TensorDataset(input, target)
train_loader = DataLoader(train_dataset, batch_size=1000, shuffle=True)

num_epochs = 1000000

for epoch in range(num_epochs):
    for step, (data, target) in enumerate(train_loader):
        # print(data)
        # print(target)
        # break
        output = model(data)
        loss = loss_fuc(output, target)

        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    print('Epoch:', epoch + 1, ' ', 'Loss: ', loss.item())

    if loss.item() < 0.0001:
        break

torch.save(model, 'net.pkl')  # 保存整个网络
torch.save(model.state_dict(), 'net_params.pkl')   # 只保存网络中的参数 (速度快, 占内存少)

三、测试神经网络

test.py

import torch
import matplotlib.pyplot as plt

net2 = torch.load('net.pkl')

# 输入参数，[Tg, theta, T, alpha]
Tg = 0.35
theta = 0.05
alpha = 0.08


gamma = round(0.9 + (0.05 - theta) / (0.3 + 6 * theta), 2)
eta1 = round(0.02 + (0.05 - theta) / (4 + 32 * theta), 3)
eta2 = round(1 + (0.05 - theta) / (0.06 + 1.6 * theta), 2)

plt.subplot()

x = []
target = []
prediction = []

for i in range(600):
    T = i * 0.01
    x.append(round(i * 0.01, 2))

    if T <= 0.1:
        target.append([round(10 * (eta2 - 0.45) * alpha * T + 0.45 * alpha, 5)])
    elif T <= Tg:
        target.append([round(eta2 * alpha, 5)])
    elif T <= 5 * Tg:
        target.append([round(pow((Tg / T), gamma) * eta2 * alpha, 5)])
    else:
        target.append([round((eta2 * pow(0.2, gamma) - eta1 * (T - 5 * Tg)) * alpha, 5)])

    input = torch.Tensor([Tg, theta, T, alpha])
    input = torch.tensor(input, dtype=torch.float32)

    output = net2(input)
    output = torch.tensor(output, dtype=torch.float32)
    output = output.numpy()
    prediction.append(output)


plt.plot(x, target, 'r-', lw=2)
plt.plot(x, prediction, 'b-', lw=2)
plt.show()

结果：