from model.lstm import LSTM
from config import args
import torch
import time
import pandas as pd
from utils.pre_process import load_data, to_transpose_lstm
import torch.utils.data as Data
import torch.nn as nn
import matplotlib.pyplot as plt
import numpy as np
from sklearn.metrics import accuracy_score, confusion_matrix
import seaborn as sns
# 调制信号类别
classes = {b'QAM16': 0, b'QAM64': 1, b'8PSK': 2, b'WBFM': 3, b'BPSK': 4,
b'CPFSK': 5, b'AM-DSB': 6, b'GFSK': 7, b'PAM4': 8, b'QPSK': 9, b'AM-SSB': 10}
# device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
device = torch.device('cpu')
def train_valid(model, args, train_dataloader, valid_dataloader, criterion, optimizer):
train_loss_all = []
train_acc_all = []
valid_loss_all = []
valid_acc_all = []
best_acc = 0.0
for epoch in range(args.epochs):
print('Epoch {}/{}'.format(epoch, args.epochs - 1))
print('-' * 50)
train_loss = 0.0
train_corrects = 0.0
train_num = 0
valid_loss = 0.0
valid_corrects = 0.0
valid_num = 0
start_time = time.time()
model.train() # 训练
for batch_idx, (input, target) in enumerate(train_dataloader):
output = model(input)
pre_lab = torch.argmax(output, 1)
loss = criterion(output, target)
optimizer.zero_grad()
loss.backward()
optimizer.step()
train_loss += loss.item() * input.size(0)
train_corrects += torch.sum(pre_lab == target)
train_num += input.size(0)
training_time = time.time() - start_time
model.eval() # 验证
with torch.no_grad():
for batch_idx, (input, target) in enumerate(valid_dataloader):
output = model(input)
pre_lab = torch.argmax(output, 1)
loss = criterion(output, target)
valid_loss += loss.item() * input.size(0)
valid_corrects += torch.sum(pre_lab == target)
valid_num += input.size(0)
train_loss_all.append(train_loss/train_num)
train_acc_all.append(train_corrects.item() / train_num)
valid_loss_all.append(valid_loss/valid_num)
valid_acc_all.append(valid_corrects.item() / valid_num)
print('{} Train Loss: {:.4f} Train Acc: {:.4f}'.format(
epoch, train_loss_all[-1], train_acc_all[-1]))
print('{} Val Loss: {:.4f} val Acc: {:.4f}'.format(
epoch, valid_loss_all[-1], valid_acc_all[-1]))
# 拷贝模型最高精度下的参数
if valid_acc_all[-1] >= best_acc:
best_acc = valid_acc_all[-1]
#######################保存模型#####################################
torch.save(model.state_dict(), "../MulScale_CLDNN-25.pkl",
_use_new_zipfile_serialization=False)
# torch.save(state_dict, "xxx.pth", _use_new_zipfile_serialization=False) #1.6之后pytorch默认使用zip文件格式来保存权重文件,导致这些权重文件无法直接被1.5及以下的pytorch加载。
####################################################################
# 使用最好模型的参数
#########################################################################
model.load_state_dict(torch.load('../MulScale_CLDNN-25.pkl')) # 加载模型
########################################################################
train_process = pd.DataFrame(
data={"epoch": range(args.epochs + 1), # 必须加1,才能和下面的train_loss_all等长度一致
"train_loss_all": train_loss_all,
"val_loss_all": valid_loss_all,
"train_acc_all": train_acc_all,
"val_acc_all": valid_acc_all})
return model, train_process
if __name__ == '__main__':
(mods, snrs, lbl), (X_train, Y_train), (X_val, Y_val), (X_test, Y_test), (train_idx, val_idx, test_idx) = load_data()
in_shp = list(X_train.shape[1])
classes = mods
train_data = Data.TensorDataset(X_train, X_test)
valid_data = Data.TensorDataset(Y_train, Y_test)
train_dataloader = Data.DataLoader(dataset=train_data, #使用的训练集
batch_size=args.batch_size, #批处理 样本大小
shuffle=True, #每次迭代前打乱数据
num_workers=1, #使用一个进程 只能开一个
)
valid_dataloader = Data.DataLoader(dataset=valid_data, #使用的训练集
batch_size=args.batch_size, #批处理 样本大小
shuffle=True, #每次迭代前打乱数据
num_workers=1, #使用一个进程 只能开一个
)
# 输出网络结构
model = LSTM()
mdoel = model.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.003)
model, train_process = train_valid(model, args, train_dataloader, valid_dataloader, criterion, optimizer)
# 可视化学习率
plt.figure(1)
plt.plot(train_process.epoch, train_process.lr_list)
plt.xlabel("epoch")
plt.ylabel("lr")
plt.title("learning rate")
# 可视化模型训练过程中的损失函数
plt.figure(figsize=(12, 4))
plt.subplot(1, 2, 1)
plt.plot(train_process.epoch, train_process.train_loss_all,
"ro-", label="Train loss")
plt.plot(train_process.epoch, train_process.val_loss_all,
"bs-", label="Val loss")
plt.legend()
plt.xlabel("epoch")
plt.ylabel("Loss")
plt.subplot(1, 2, 2)
plt.plot(train_process.epoch, train_process.train_acc_all,
"ro-", label="Train acc")
plt.plot(train_process.epoch, train_process.val_acc_all,
"bs-", label="Val acc")
plt.xlabel("epoch")
plt.ylabel("acc")
plt.legend()
plt.show()
# 如果测试集过大,难以计算整体模型识别率和画出混淆矩阵
(mods, snrs, lbl), (X_train, Y_train), (X_val, Y_val), (X_test, Y_test), (train_idx, val_idx, test_idx) = load_data()
print("X_test.shape:", X_test.shape)
print('-' * 50)
# 测试集准确度随信噪比的变化曲线
acc = {}
model.eval()
for snr in snrs:
test_SNRs = map(lambda x: lbl[x], test_idx)
test_SNRs = list(test_SNRs)
test_X_i = X_test[np.where(np.array(test_SNRs) == snr)]
test_Y_i = Y_test[np.where(np.array(test_SNRs) == snr)]
output = model(test_X_i)
pre_lab = torch.argmax(output, 1).cpu() # 得到最大值的序号索引
acc[snr] = accuracy_score(test_Y_i, pre_lab)
print(acc[snr])
plt.figure(3)
plt.plot(snrs, list(map(lambda x: acc[x], snrs)))
plt.xlabel("Signal to Noise Ratio")
plt.ylabel("Classification Accuracy")
plt.title("LSTM Classification Accuracy on RadioML 2016.10 Alpha")
plt.yticks(np.linspace(0, 1, 11))
acc_total = {}
output = model(X_test)
pre_lab = torch.argmax(output, 1).cpu()
acc_total = accuracy_score(Y_test, pre_lab)
print("在测试集上的预测精度为:", acc_total)
# 计算混淆矩阵并可视化
plt.figure(4)
conf_mat = confusion_matrix(Y_test, pre_lab)
df_cm = pd.DataFrame(conf_mat, index=classes,
columns=classes)
heatmap = sns.heatmap(df_cm, annot=True, fmt="d", cmap="YlGnBu")
heatmap.yaxis.set_ticklabels(
heatmap.yaxis.get_ticklabels(), rotation=0, ha='right')
heatmap.xaxis.set_ticklabels(
heatmap.xaxis.get_ticklabels(), rotation=45, ha='right')
plt.ylabel('True label')
plt.xlabel('Predicted label')
plt.show()
训练和测试代码
最新推荐文章于 2024-08-23 22:47:58 发布
本文详细介绍了如何使用Python进行深度学习模型的训练和测试。从数据预处理、模型构建、损失函数选择到训练过程优化,以及验证集评估和测试集预测,全方位解析了机器学习项目的实战步骤。
摘要由CSDN通过智能技术生成