代码非常简单。

import torch 
import torch.nn as nn
from torch.autograd import Variable
from scipy.io.wavfile import write
#need install pydub module
#pip install pydub
import numpy as np
import pydub 
from scipy import signal
import IPython
import matplotlib.pylab as plt
from mpl_toolkits.mplot3d import Axes3D
# For running on GPU
#device = torch.device("cuda")# choose your device
device = torch.device("cpu")
a = torch.rand(5, 5, device=device)# change by either using the device argument
a = a.to(device)# or by .to()

Make data

fs = 512
x = np.linspace(0, 20*np.pi * (1-1/(10*fs)), fs*10)
y_sin = 0.5*np.sin(x)
plt.plot(x, y_sin)
plt.xlabel('Angle [rad]')
plt.ylabel('sin(x)')
plt.axis('tight')
plt.show()

y_triangle = 0.5*signal.sawtooth(x, 0.5)
plt.plot(x, y_triangle)
plt.xlabel('Phase [rad]')
plt.ylabel('triangle(x)')
plt.axis('tight')
plt.show()

y_saw = 0.5*signal.sawtooth(x, 1)plt.plot(x, y_saw)plt.xlabel('Phase [rad]')plt.ylabel('sawtooth(x)')plt.axis('tight')plt.show()

Add Gaussian Noise

Add noise

# Add guassian noise
y_sin_n = y_sin + 0.1*np.random.normal(size=len(x))
y_triangle_n = y_triangle + 0.1*np.random.normal(size=len(x))
y_saw_n = y_saw + 0.1*np.random.normal(size=len(x))plt.plot(x, y_sin_n)
plt.xlabel('Angle [rad]')
plt.ylabel('sin(x) + noise')
plt.axis('tight')
plt.show()

plt.plot(x, y_triangle_n)
plt.xlabel('Phase [rad]')
plt.ylabel('triangle(x) + noise')
plt.axis('tight')
plt.show()

plt.plot(x, y_saw_n)
plt.xlabel('Phase [rad]')
plt.ylabel('sawtooth(x) + noise')
plt.axis('tight')
plt.show()

Creating Dataset

def give_part_of_data(x, y, n_samples=10000, sample_size=100) :data_inp = np.zeros((n_samples, sample_size))data_out = np.zeros((n_samples, sample_size))for i in range(n_samples):random_offset = np.random.randint(0, len(x) - sample_size)sample_inp = x[random_offset:random_offset+sample_size]sample_out = y[random_offset:random_offset+sample_size]data_inp[i, :] = sample_inpdata_out[i, :] = sample_outreturn data_inp, data_out
# Train, Validationa, and Test
sin_train_in, sin_train_out = give_part_of_data(y_sin_n[0:int(7/10 * len(x))], y_sin[0:int(7/10 * len(x))], 2000, int(len(x)/6))
tri_train_in, tri_train_out = give_part_of_data(y_triangle_n[0:int(7/10 * len(x))], y_triangle[0:int(7/10 * len(x))], 2000, int(len(x)/6))
saw_train_in, saw_train_out = give_part_of_data(y_saw_n[0:int(7/10 * len(x))], y_saw[0:int(7/10 * len(x))], 2000, int(len(x)/6))sin_val_in, sin_val_out = y_sin_n[int(7/10 * len(x)):int(8/10 * len(x))], y_sin[int(7/10 * len(x)):int(8/10 * len(x))]
tri_val_in, tri_val_out = y_triangle_n[int(7/10 * len(x)):int(8/10 * len(x))], y_triangle[int(7/10 * len(x)):int(8/10 * len(x))]
saw_val_in, saw_val_out = y_saw_n[int(7/10 * len(x)):int(8/10 * len(x))], y_saw[int(7/10 * len(x)):int(8/10 * len(x))]sin_test_in, sin_test_out = y_sin_n[int(8/10 * len(x)):int(10/10 * len(x))], y_sin[int(8/10 * len(x)):int(10/10 * len(x))]
tri_test_in, tri_test_out = y_triangle_n[int(8/10 * len(x)):int(10/10 * len(x))], y_triangle[int(8/10 * len(x)):int(10/10 * len(x))]
saw_test_in, saw_test_out = y_saw_n[int(8/10 * len(x)):int(10/10 * len(x))], y_saw[int(8/10 * len(x)):int(10/10 * len(x))]
plt.plot(range(853), sin_train_in[3])
plt.plot(range(853), sin_train_out[3])plt.xlabel('Phase [rad]')
plt.ylabel('sin(x) + noise')
plt.axis('tight')
plt.show()

RNN + Sin

# RNN model
input_dim = 1
hidden_size_1 = 60
hidden_size_2 = 60
output_size = 1class CustomRNN(nn.Module):def __init__(self, input_size, hidden_size_1, hidden_size_2, output_size):super(CustomRNN, self).__init__()self.rnn = nn.RNN(input_size=input_size, hidden_size=hidden_size_1, batch_first=True)self.linear = nn.Linear(hidden_size_1, hidden_size_2, )self.act = nn.Tanh()self.linear = nn.Linear(hidden_size_2, output_size, )self.act = nn.Tanh()def forward(self, x):pred, hidden = self.rnn(x, None)pred = self.act(self.linear(pred)).view(pred.data.shape[0], -1, 1)return predmodel = CustomRNN(input_dim, hidden_size_1, hidden_size_2, output_size)
model = model.to(device)
optimizer = torch.optim.Adam(model.parameters())
loss_func = nn.MSELoss()lr = 1e-2for t in range(1000):inp = torch.Tensor(sin_train_in[..., np.newaxis] )inp.requires_grad = Trueinp = inp.to(device)out = torch.Tensor(sin_train_out[..., np.newaxis])out = out.to(device)pred = model(inp)optimizer.zero_grad()loss = loss_func(pred, out)if t%20==0:print(t, loss.data.item())lr = lr / 1.0001optimizer.param_groups[0]['lr'] = lrloss.backward()optimizer.step()
test_in = sin_test_in
inp = torch.Tensor(test_in[np.newaxis, ... , np.newaxis] )
inp = inp.to(device)
pred = model(inp).cpu().detach().numpy()
plt.plot(range(len(sin_test_in)), test_in)
plt.plot(range(len(sin_test_in)), pred[0, :,0])plt.showorginal_SNR = np.sum(np.abs(sin_test_out)**2) / np.sum(np.abs(sin_test_in - sin_test_out)**2)
orginal_SNR_db = 10*np.log(orginal_SNR)/np.log(10)
print('Original SNR : ', orginal_SNR)
print('Original SNR DB : ', orginal_SNR_db)network_SNR = np.sum(np.abs(sin_test_out)**2) / np.sum(np.abs(pred[0, :,0] - sin_test_out)**2)
network_SNR_db = 10*np.log(network_SNR)/np.log(10)
print('Network SNR : ', network_SNR)
print('Network SNR DB : ', network_SNR_db)
Original SNR :  12.951857235597608
Original SNR DB :  11.123320486750668
Network SNR :  107.29848229242438
Network SNR DB :  20.305935790331755

RNN + Triangular

# RNN model
input_dim = 1
hidden_size_1 = 60
hidden_size_2 = 60
output_size = 1class CustomRNN(nn.Module):def __init__(self, input_size, hidden_size_1, hidden_size_2, output_size):super(CustomRNN, self).__init__()self.rnn = nn.RNN(input_size=input_size, hidden_size=hidden_size_1, batch_first=True)self.linear = nn.Linear(hidden_size_1, hidden_size_2, )self.act = nn.Tanh()self.linear = nn.Linear(hidden_size_2, output_size, )self.act = nn.Tanh()def forward(self, x):pred, hidden = self.rnn(x, None)pred = self.act(self.linear(pred)).view(pred.data.shape[0], -1, 1)return predmodel = CustomRNN(input_dim, hidden_size_1, hidden_size_2, output_size)
model = model.to(device)
optimizer = torch.optim.Adam(model.parameters())
loss_func = nn.MSELoss()lr = 1e-2for t in range(1000):inp = torch.Tensor(tri_train_in[..., np.newaxis] )inp.requires_grad = Trueinp = inp.to(device)out = torch.Tensor(tri_train_out[..., np.newaxis])out = out.to(device)pred = model(inp)optimizer.zero_grad()loss = loss_func(pred, out)if t%20==0:print(t, loss.data.item())lr = lr / 1.0001optimizer.param_groups[0]['lr'] = lrloss.backward()optimizer.step()
test_in = tri_test_in
inp = torch.Tensor(test_in[np.newaxis, ... , np.newaxis] )
inp = inp.to(device)
pred = model(inp).cpu().detach().numpy()
plt.plot(range(len(tri_test_in)), test_in)
plt.plot(range(len(tri_test_in)), pred[0, :,0])plt.showorginal_SNR = np.sum(np.abs(tri_test_out)**2) / np.sum(np.abs(tri_test_in - tri_test_out)**2)
orginal_SNR_db = 10*np.log(orginal_SNR)/np.log(10)
print('Original SNR : ', orginal_SNR)
print('Original SNR DB : ', orginal_SNR_db)network_SNR = np.sum(np.abs(tri_test_out)**2) / np.sum(np.abs(pred[0, :,0] - tri_test_out)**2)
network_SNR_db = 10*np.log(network_SNR)/np.log(10)
print('Network SNR : ', network_SNR)
print('Network SNR DB : ', network_SNR_db)
Original SNR :  9.06282337035853
Original SNR DB :  9.572635159053185
Network SNR :  46.622532666082044
Network SNR DB :  16.685958619136

RNN + Sawtooth

# RNN model
input_dim = 1
hidden_size_1 = 60
hidden_size_2 = 60
output_size = 1class CustomRNN(nn.Module):def __init__(self, input_size, hidden_size_1, hidden_size_2, output_size):super(CustomRNN, self).__init__()self.rnn = nn.RNN(input_size=input_size, hidden_size=hidden_size_1, batch_first=True)self.linear = nn.Linear(hidden_size_1, hidden_size_2, )self.act = nn.Tanh()self.linear = nn.Linear(hidden_size_2, output_size, )self.act = nn.Tanh()def forward(self, x):pred, hidden = self.rnn(x, None)pred = self.act(self.linear(pred)).view(pred.data.shape[0], -1, 1)return predmodel = CustomRNN(input_dim, hidden_size_1, hidden_size_2, output_size)
model = model.to(device)
optimizer = torch.optim.Adam(model.parameters())
loss_func = nn.MSELoss()lr = 1e-2for t in range(1000):inp = torch.Tensor(tri_train_in[..., np.newaxis] )inp.requires_grad = Trueinp = inp.to(device)out = torch.Tensor(tri_train_out[..., np.newaxis])out = out.to(device)pred = model(inp)optimizer.zero_grad()loss = loss_func(pred, out)if t%20==0:print(t, loss.data.item())lr = lr / 1.0001optimizer.param_groups[0]['lr'] = lrloss.backward()optimizer.step()
test_in = saw_test_in
inp = torch.Tensor(test_in[np.newaxis, ... , np.newaxis] )
inp = inp.to(device)
pred = model(inp).cpu().detach().numpy()
plt.plot(range(len(saw_test_in)), test_in)
plt.plot(range(len(saw_test_in)), pred[0, :,0])plt.showorginal_SNR = np.sum(np.abs(saw_test_out)**2) / np.sum(np.abs(saw_test_in - saw_test_out)**2)
orginal_SNR_db = 10*np.log(orginal_SNR)/np.log(10)
print('Original SNR : ', orginal_SNR)
print('Original SNR DB : ', orginal_SNR_db)network_SNR = np.sum(np.abs(saw_test_out)**2) / np.sum(np.abs(pred[0, :,0] - saw_test_out)**2)
network_SNR_db = 10*np.log(network_SNR)/np.log(10)
print('Network SNR : ', network_SNR)
print('Network SNR DB : ', network_SNR_db)
Original SNR :  8.918716305325825
Original SNR DB :  9.50302349708762
Network SNR :  26.97065260659425
Network SNR DB :  14.308914551667852

知乎学术咨询：
https://www.zhihu.com/consult/people/792359672131756032?isMe=1

工学博士，担任《Mechanical System and Signal Processing》《中国电机工程学报》《控制与决策》等期刊审稿专家，擅长领域：现代信号处理，机器学习，深度学习，数字孪生，时间序列分析，设备缺陷检测、设备异常检测、设备智能故障诊断与健康管理PHM等。