<think>好的,我现在需要帮助用户完成一个基于PyTorch的调制识别算法,使用ResNet50和迁移学习。用户的要求挺详细的,我得一步步来梳理。
首先,用户提到的数据集是RML2016.10a_dict.pkl,这个数据集我之前接触过,是无线信号调制识别的常用数据集,通常包含不同信噪比下的信号样本。数据预处理部分可能需要将信号转换为时频图,比如使用短时傅里叶变换(STFT)生成图像,然后作为ResNet50的输入。
接下来是迁移学习。用户要求冻结除全连接层和部分卷积层之外的层。ResNet50的模型结构里,通常前面的卷积层用于提取低级特征,后面的层处理更高级的特征。所以可能需要冻结前面的某些块,比如前三个或四个块,然后允许后面的卷积层和全连接层进行微调。但具体冻结哪些层可能需要根据实际情况调整,可能需要查看ResNet50的结构来确定。
然后是设置随机种子,这涉及到PyTorch的随机种子、numpy的随机种子,还有CUDA的确定性设置,确保实验的可重复性。需要用到torch.manual_seed,np.random.seed,以及设置cudnn.deterministic和cudnn.benchmark为False。
模型结构方面,ResNet50的原始全连接层输出是1000类,这里需要替换为数据集的类别数。用户还提到要添加Dropout层和L2正则化。Dropout可以在全连接层之前添加,而L2正则化可以通过在优化器中设置weight_decay参数来实现。
早停法(Early Stopping)需要在训练过程中监控验证集的损失,如果一段时间内没有改善,就停止训练。这需要记录每个epoch的验证损失,并在达到耐心次数时终止训练。
交叉验证部分,用户要求使用交叉验证。但通常PyTorch中实现交叉验证需要手动分割数据集,这里可能需要将数据集分成K折,比如5折,然后依次训练和验证。不过,考虑到计算资源,可能用户希望的是简单的K折交叉验证,或者只是单次的验证集划分。需要确认,但根据用户给出的示例代码,可能采用简单的训练-验证分割。
可视化部分需要绘制损失曲线、准确率曲线、混淆矩阵、各类信号的SNR-准确率曲线、ROC曲线和AUC值。这些都需要在训练过程中收集相应的数据,比如每个epoch的训练和验证损失与准确率,测试阶段的真实标签和预测结果,不同SNR下的准确率,以及每个类别的ROC曲线。
保存结果到output1文件夹,需要确保在代码中创建该目录,并将图片和模型等保存进去。
关于代码结构,可能需要以下步骤:
1. 导入必要的库,包括PyTorch、numpy、matplotlib、pickle、sklearn等。
2. 设置随机种子,确保可重复性。
3. 加载数据集,可能需要处理.pkl文件,提取信号数据和标签。
4. 数据预处理,将信号转换为时频图,可能需要调整形状和归一化。
5. 定义数据集类,继承Dataset,处理数据加载。
6. 定义模型,修改ResNet50,冻结部分层,添加Dropout,修改全连接层。
7. 定义训练和验证的函数,包括损失计算、优化器(带L2正则化)、早停法。
8. 进行交叉验证循环,比如5折交叉验证。
9. 在训练过程中收集损失和准确率数据。
10. 训练完成后,进行测试,生成混淆矩阵、SNR-准确率曲线、ROC曲线等。
11. 保存模型和绘图结果到output1文件夹。
需要注意的细节:
- 数据预处理部分,信号转换为时频图可能需要使用librosa的stft函数,或者PyTorch自己的方法。需要确保生成的时频图是适合ResNet50输入的3通道图像,可能通过堆叠相同的单通道图像来模拟RGB。
- 类别标签需要从数据集中正确提取,并转换为数字标签。
- 在冻结层时,需要遍历模型的参数,设置requires_grad为False,对需要训练的部分设置为True。
- 交叉验证的实现可能需要将数据集划分为K个子集,依次选择其中一个作为验证集,其余作为训练集。但处理时间可能较长,用户是否希望完整的K折交叉验证?或者只是单次划分?示例代码可能采用单次划分,但用户提到交叉验证,可能需要调整。
- 早停法需要跟踪验证损失,并在连续多个epoch没有改善时停止。
- L2正则化通过优化器的weight_decay参数实现,需要在定义优化器时设置。
- 混淆矩阵需要使用sklearn的confusion_matrix,然后绘制热力图。
- SNR-准确率曲线需要按每个SNR值计算该SNR下的准确率,可能需要将测试数据按SNR分组,分别计算。
- ROC曲线和AUC值需要处理多分类情况,可能使用OneVsRest策略,计算每个类别的ROC曲线,并取宏平均或其他方式。
可能遇到的问题:
- 数据集的结构:需要确认RML2016.10a_dict.pkl的数据结构,通常该数据集是一个字典,键是调制类型和SNR,值是对应的信号数据。每个样本可能是复数时间序列,需要转换为时频图。
- 时频图的生成是否正确,可能需要进行归一化,调整尺寸以适应ResNet50的输入(通常224x224)。
- 内存问题,如果数据集很大,生成时频图可能需要大量内存,可能需要逐批处理。
- 类别不平衡问题,不同调制类型在不同SNR下的样本数量可能不同,可能需要统计并确保训练时不会偏向多数类。
- 交叉验证的实现是否与早停法兼容,因为在每个fold中可能需要重新初始化模型,并重新进行早停判断。
现在需要将上述思考转化为代码结构,确保各部分正确衔接,并处理可能的异常情况,比如目录创建、数据分割时的随机性控制等。</think>以下是根据您需求构建的完整PyTorch实现代码,代码分为数据预处理、模型定义、训练流程和可视化四个部分:
```python
import os
import pickle
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader, Subset
from torchvision import transforms
from sklearn.model_selection import KFold, train_test_split
from sklearn.metrics import confusion_matrix, roc_curve, auc
from sklearn.preprocessing import label_binarize
from torchvision.models import resnet50
from scipy.signal import stft
# 设置随机种子确保可重复性
seed = 42
torch.manual_seed(seed)
np.random.seed(seed)
os.makedirs('output1', exist_ok=True)
#%% 数据预处理
def load_data(path):
with open(path, 'rb') as f:
data = pickle.load(f, encoding='bytes')
return data
def signal_to_spectrogram(signal, nperseg=128):
f, t, Zxx = stft(signal, nperseg=nperseg)
return np.abs(Zxx)
class SignalDataset(Dataset):
def __init__(self, data_dict):
self.samples = []
self.labels = []
self.snrs = []
self.classes = sorted(list(set(k[0] for k in data_dict.keys())))
for idx, (key, val) in enumerate(data_dict.items()):
mod_type, snr = key
for signal in val:
spec = signal_to_spectrogram(signal)
spec = (spec - spec.min()) / (spec.max() - spec.min()) # 归一化
spec = np.stack([spec]*3, axis=0) # 转为3通道
self.samples.append(spec)
self.labels.append(self.classes.index(mod_type))
self.snrs.append(snr)
def __len__(self):
return len(self.samples)
def __getitem__(self, idx):
return (
torch.FloatTensor(self.samples[idx]),
self.labels[idx],
self.snrs[idx]
)
#%% 模型定义
class ModifiedResNet50(nn.Module):
def __init__(self, num_classes, dropout_prob=0.5):
super().__init__()
self.base_model = resnet50(pretrained=True)
# 冻结前4个卷积块
for param in self.base_model.parameters():
param.requires_grad = False
for param in self.base_model.layer4.parameters():
param.requires_grad = True
# 修改全连接层
self.base_model.fc = nn.Sequential(
nn.Dropout(dropout_prob),
nn.Linear(2048, num_classes)
)
def forward(self, x):
return self.base_model(x)
#%% 训练配置
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
BATCH_SIZE = 64
EPOCHS = 100
PATIENCE = 7
K_FOLDS = 5
# 加载数据
data_dict = load_data('/workspace/rrr/RML2016.10a_dict.pkl')
full_dataset = SignalDataset(data_dict)
# 交叉验证
kfold = KFold(n_splits=K_FOLDS, shuffle=True, random_state=seed)
results = {}
for fold, (train_idx, test_idx) in enumerate(kfold.split(full_dataset)):
print(f'Fold {fold+1}')
# 划分训练验证集
train_idx, val_idx = train_test_split(train_idx, test_size=0.2, random_state=seed)
# 数据加载
train_loader = DataLoader(Subset(full_dataset, train_idx), batch_size=BATCH_SIZE, shuffle=True)
val_loader = DataLoader(Subset(full_dataset, val_idx), batch_size=BATCH_SIZE)
test_loader = DataLoader(Subset(full_dataset, test_idx), batch_size=BATCH_SIZE)
# 初始化模型
model = ModifiedResNet50(num_classes=len(full_dataset.classes)).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()),
lr=1e-4, weight_decay=1e-4) # L2正则化
# 早停法
best_val_loss = np.inf
epochs_no_improve = 0
history = {'train_loss':[], 'val_loss':[], 'train_acc':[], 'val_acc':[]}
# 训练循环
for epoch in range(EPOCHS):
# 训练阶段
model.train()
running_loss = 0.0
correct = 0
total = 0
for inputs, labels, _ in train_loader:
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
train_loss = running_loss/len(train_loader)
train_acc = correct/total
# 验证阶段
model.eval()
val_loss = 0.0
correct = 0
total = 0
with torch.no_grad():
for inputs, labels, _ in val_loader:
inputs, labels = inputs.to(device), labels.to(device)
outputs = model(inputs)
loss = criterion(outputs, labels)
val_loss += loss.item()
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
val_loss = val_loss/len(val_loader)
val_acc = correct/total
# 记录历史
history['train_loss'].append(train_loss)
history['val_loss'].append(val_loss)
history['train_acc'].append(train_acc)
history['val_acc'].append(val_acc)
# 早停判断
if val_loss < best_val_loss:
best_val_loss = val_loss
epochs_no_improve = 0
torch.save(model.state_dict(), f'output1/best_model_fold{fold}.pth')
else:
epochs_no_improve += 1
if epochs_no_improve == PATIENCE:
print(f'Early stopping at epoch {epoch+1}')
break
# 保存训练曲线
plt.figure(figsize=(12,5))
plt.subplot(121)
plt.plot(history['train_loss'], label='Train')
plt.plot(history['val_loss'], label='Validation')
plt.title('Loss Curve')
plt.legend()
plt.subplot(122)
plt.plot(history['train_acc'], label='Train')
plt.plot(history['val_acc'], label='Validation')
plt.title('Accuracy Curve')
plt.legend()
plt.savefig(f'output1/training_curves_fold{fold}.png')
plt.close()
#%% 测试与可视化
model.load_state_dict(torch.load(f'output1/best_model_fold{fold}.pth'))
model.eval()
all_preds = []
all_labels = []
all_snrs = []
with torch.no_grad():
for inputs, labels, snrs in test_loader:
inputs = inputs.to(device)
outputs = model(inputs)
_, preds = torch.max(outputs, 1)
all_preds.extend(preds.cpu().numpy())
all_labels.extend(labels.numpy())
all_snrs.extend(snrs.numpy())
# 混淆矩阵
cm = confusion_matrix(all_labels, all_preds)
plt.figure(figsize=(10,8))
plt.imshow(cm, interpolation='nearest', cmap=plt.cm.Blues)
plt.title('Confusion Matrix')
plt.colorbar()
plt.xticks(np.arange(len(full_dataset.classes)), full_dataset.classes, rotation=45)
plt.yticks(np.arange(len(full_dataset.classes)), full_dataset.classes)
plt.tight_layout()
plt.savefig(f'output1/cm_fold{fold}.png')
plt.close()
# SNR-准确率曲线
snr_values = sorted(list(set(all_snrs)))
acc_per_snr = []
for snr in snr_values:
indices = [i for i, s in enumerate(all_snrs) if s == snr]
acc = np.mean([all_preds[i] == all_labels[i] for i in indices])
acc_per_snr.append(acc)
plt.figure()
plt.plot(snr_values, acc_per_snr, marker='o')
plt.xlabel('SNR (dB)')
plt.ylabel('Accuracy')
plt.title('SNR-Accuracy Curve')
plt.grid()
plt.savefig(f'output1/snr_acc_fold{fold}.png')
plt.close()
# ROC曲线和AUC
y_test = label_binarize(all_labels, classes=range(len(full_dataset.classes)))
y_score = np.zeros_like(y_test)
with torch.no_grad():
for inputs, _, _ in test_loader:
inputs = inputs.to(device)
outputs = model(inputs)
y_score = np.vstack([y_score, outputs.cpu().numpy()])
fpr = dict()
tpr = dict()
roc_auc = dict()
for i in range(len(full_dataset.classes)):
fpr[i], tpr[i], _ = roc_curve(y_test[:,i], y_score[:,i])
roc_auc[i] = auc(fpr[i], tpr[i])
plt.figure()
for i in range(len(full_dataset.classes)):
plt.plot(fpr[i], tpr[i], label=f'{full_dataset.classes[i]} (AUC={roc_auc[i]:.2f})')
plt.plot([0,1],[0,1],'k--')
plt.xlabel('False Positive Rate')
plt.ylabel('True Positive Rate')
plt.title('ROC Curve')
plt.legend()
plt.savefig(f'output1/roc_fold{fold}.png')
plt.close()
#%% 结果汇总
print("所有结果已保存在output1文件夹中")
```
### 代码说明:
1. **数据预处理**:
- 使用STFT将信号转换为时频图
- 进行归一化和三通道处理
- 创建自定义Dataset类
2. **模型架构**:
- 基于预训练的ResNet50
- 冻结前4个卷积块
- 添加Dropout层和L2正则化
3. **训练特性**:
- 5折交叉验证
- 早停法(耐心值7)
- Adam优化器带权重衰减(L2正则化)
- 学习率1e-4
4. **可视化输出**:
- 每折保存训练曲线
- 混淆矩阵
- SNR-准确率曲线
- 多分类ROC曲线和AUC值
5. **结果保存**:
- 所有输出自动保存在output1文件夹
- 包含模型权重和各类分析图表
### 使用建议:
1. 确保PyTorch环境已安装scipy、matplotlib等依赖库
2. 根据GPU显存调整batch_size
3. 可通过修改K_FOLDS值调整交叉验证次数
4. 最终结果包含多个分析图表,建议使用Jupyter Notebook查看