import numpy as np import torch import torch.nn as nn import torch.optim as optim form torch.utils.data import Dataset, DataLoader # 生成信号 def generate_signal(freq, sample_rate, length, phase, snr): t = np.arange(length) / sample_rate signal = np.cos(2 * np.pi * freq * t + phase) + 1j * np.sin(2 * np.pi * freq * t + phase) noise_power = 10 ** (-snr/10) noise = np.sqrt(noise_power) * (np.random.randn(length) + 1j * np.random.randn(length) ) return signal + noise class SignalDataset(Dataset): def _ _init_ _(self, num_sample, sample_rata, length, freqs, snr_range): self.num_samples = num_sample self.samples_rate = sample_rate self.length = length self.freqs = freqs self.snr_range = snr_range self.data, self.labels = self.generate_data() def generate_data(self): data = [] labels = [] for i in range(self.num_samples): freq = np.random.choice(self.freqs) phase = np.random.uniform( 0, 2 * np.pi) snr = np.random.uniform(self.snr_range[0], self.snr_range[1]) signal = generate_signal(freq, self.sample_rate, self.length, phase, snr) data.append(signal) labels.append(self.freqs.index(freq)) return np.array(data), np.array(labels) def _ _len_ _(self, idx): return self.num_samples def _ _getitem_ _(self, idx): signal = self.data[idx] label = self.labels[idx] return torch.tensor(signal.real, dtype=torch.float32), torch.tensor(signal.imag, dtype=torch.float32), torch.tensor(label, dtype=torch.long) # 数据生成参数设置 num_samples = 3000 sample_rate = 640 length = 256 freqs = [10, 50, 90] snr_range = [5, 10] # 生成数据集 dataset = SignalDataset(num_samples, sample_rate, length, freqs, snr_range) train_size = int(0.8 * num_samples) test_size = num_samples - train_size train_dataset, test_dataset = torch.utils.data.random_split(dataset, [train_size, test_size]) # 创建数据集加载器 train_loader = DataLoader(train_dataset, batch_size=32, shuffle= True) test_loader = DataLoader(test_dataset, batch_size=32, shuffle= False) # 模型构建 class LSTMClassifier(nn.Moule): def _ _init_ _(self, input_size, hidden_size, num_layers, num_classes): super(LSTMClassifier, self)._ _init_ _() self.hidden_size = hidden_size self.num_layers =num_layers self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True) self.fc = nn.Linear(hidden_size, num_classes) def forward(self, x): h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device) c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device) out, _ = self.lstm(x, (h0, c0)) out = self.fc(out[:, -1, :]) return out # 训练函数 def train(model, dataloader, criterion, optimizer, device): model.train() running_loss = 0.0 for real, imag, labels in dataloader: inputs = torch.stack([real, imag], dim=-1).to(device) labels = labels.to(device) optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() running_loss += loss.item() return running_loss / len(dataloader) # 测试函数 def test(model, dataloader, criterion, device): model.eval() running_loss = 0.0 correct = 0 total = 0 with torch.no_grad(): for real, imag, labels in dataloader: inputs = torch.stack([real, imag], dim=-1).to(device) labels = labels.to(device) outputs = model(inputs) loss = criterion(outputs, labels) running_loss += loss.item() _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (prediected ==labels).sum().item() accuracy = correct /total return running_loss / len(dataloader), accuracy # 模型训练测试参数设置 input_size = 2 # 复信号的实部和虚部 hidden_size = 64 num_layers = 2 num_classes = len(freqs) # 模型构建 device = 'cude' if torch.cuda.is_available() else 'cpu' model = LSTMClassifier(input_size, hidden_size, num_layers, num_classes) model = model.to(device) # 损失函数和优化器定义 criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=0.001) # 训练和测试 num_epochs = 100 for epoch in range(num_epochs): train_loss = train(model, train_loader, criterion, optimizer, device) test_loss, test_acc = test(model, test_loader, criterion, optimizer, device) printf(f'Epoch [{epoch + 1} / {num_epochs}], Train Loss: {train_loss: .4f}, Test Loss: {test_loss: .4f}, Test Accuracy: {test_acc: .4f}) # 测试 model.eval() correct = 0 total = 0 with torch.no_grad() for real, imag, labels in test_loader: inputs = torch.stack([real, imag], dim=-1).to(device) labels = labels.to(device) outputs = model(inputs) _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (prediected ==labels).sum().item() accuracy = correct / total print(f' Test Accuracy: {accuracy: .4f}') | 
