请使用RNN+CNN复原波形

好的，下面是一个使用 RNN+CNN 复原波形的代码实现。需要注意的是，本实现中 RNN+CNN 的结构是在 RNN 的每个时间步上应用了一个 CNN，而不是将 CNN 和 RNN 相互嵌套。实现中仍然使用 PTB 心电数据集进行演示，需要先安装 wfdb、pywt 和 torch 等库：

import wfdb
import pywt
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader

# 读取数据
record = wfdb.rdrecord('ptbdb/patient001/s0010_re')
signal = record.p_signal[:,0]

# 滤波和去噪
def denoise(signal):
    # 小波去噪
    coeffs = pywt.wavedec(signal, 'db4', level=4)
    coeffs[1:] = [pywt.threshold(c, 0.1*np.max(c)) for c in coeffs[1:]]
    signal_denoised = pywt.waverec(coeffs, 'db4')
    
    # 中值滤波
    signal_filtered = np.zeros_like(signal_denoised)
    for i in range(1, len(signal_filtered)-1):
        signal_filtered[i] = np.median(signal_denoised[i-1:i+2])
    signal_filtered[0] = signal_filtered[1]
    signal_filtered[-1] = signal_filtered[-2]
    
    return signal_filtered

signal_processed = denoise(signal)

# 定义数据集
class EcgDataset(Dataset):
    def __init__(self, signal, window_size=1000, stride=100):
        self.signal = signal
        self.window_size = window_size
        self.stride = stride
        
    def __len__(self):
        return (len(self.signal) - self.window_size) // self.stride + 1
        
    def __getitem__(self, idx):
        start = idx * self.stride
        end = start + self.window_size
        x = self.signal[start:end].reshape(1, -1)
        y = self.signal[end:end+1]
        return x, y

# 定义模型
class EcgModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.rnn = nn.GRU(1, 64, batch_first=True, bidirectional=True)
        self.cnn = nn.Sequential(
            nn.Conv1d(128, 64, kernel_size=5, stride=1, padding=2),
            nn.ReLU(),
            nn.Conv1d(64, 1, kernel_size=1, stride=1),
        )
        
    def forward(self, x):
        batch_size = x.shape[0]
        x, _ = self.rnn(x)
        x = x.transpose(1, 2)
        x = self.cnn(x)
        x = x.transpose(1, 2)
        return x.view(batch_size, -1)

# 训练模型
train_dataset = EcgDataset(signal_processed[:50000])
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
model = EcgModel()
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.MSELoss()

for epoch in range(10):
    for x, y in train_loader:
        optimizer.zero_grad()
        y_pred = model(x)
        loss = criterion(y_pred, y)
        loss.backward()
        optimizer.step()
    print(f'Epoch {epoch+1}, loss={loss.item()}')

# 测试模型并复原波形
test_dataset = EcgDataset(signal_processed[50000:])
test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False)

with torch.no_grad():
    model.eval()
    reconstructed_signal = []
    for x, y in test_loader:
        x = x.transpose(1, 2)
        x = model.rnn(x)[0]
        x = x.transpose(1, 2)
        y_pred = model.cnn(x)
        reconstructed_signal.append(y_pred.item())
    reconstructed_signal = np.concatenate([signal_processed[:50000], reconstructed_signal])

# 显示结果
import matplotlib.pyplot as plt

plt.figure(figsize=(10,5))
plt.subplot(1,2,1)
plt.plot(signal[:5000])
plt.title('Raw ECG Signal')
plt.subplot(1,2,2)
plt.plot(reconstructed_signal[:5000])
plt.title('Reconstructed ECG Signal')
plt.show()

这个代码实现了对 PTB 心电数据集的预处理、使用 RNN+CNN 进行特征提取和 RNN+CNN 对数据进行复原，其中 RNN 用于提取序列信息，CNN 用于提取局部特征，最后通过线性映射将特征转换为输出值。需要注意的是，在复原波形时，需要将输入序列先通过 RNN 提取序列信息，再将输出结果输入到 CNN 中，得到复原后的序列。代码中使用均方误差作为损失函数，Adam 作为优化器，训练 10 轮。最后将复原后的波形和原始波形进行比较。