from data_process import get_data import torch from sklearn.model_selection import train_test_split from LeNet5 import LeNet5 X, y = get_data() # 获取数据【0.025,0.035】100*0.2 = 20 X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, stratify=y) # 数据拆分 print(X_train.shape) #(1075, 227, 227, 1) 0 1 2 3 --- (1075, 1, 227, 227) 0 3 1 2 X_train_tensor = torch.tensor(X_train, dtype=torch.float32).permute(0, 3, 1, 2) # 将数据转成模型要求的形式 print(X_train_tensor.shape) X_test_tensor = torch.tensor(X_test, dtype=torch.float32).permute(0, 3, 1, 2) y_train_tensor = torch.tensor(y_train, dtype=torch.int64) train_ds = torch.utils.data.TensorDataset(X_train_tensor, y_train_tensor) # 将数据转为tensordata类型 train_dl = torch.utils.data.DataLoader(train_ds, batch_size=128, shuffle=True) # 对数据进行分批及打乱操作 network = LeNet5() # 实例化得到一个leNet-5网络模型 loss_fn = torch.nn.CrossEntropyLoss() # 损失函数(交差熵) optimizer = torch.optim.SGD(network.parameters(), lr=0.01) # 优化器 # 模型训练 for epoch in range(1): for image, label in train_dl: y_pre = network(image) # 模型计算(前向传播) loss = loss_fn(y_pre, label) # 计算损失值 network.zero_grad() # 将网络中的所有梯度清零 loss.backward() # 计算梯度项(反向求导) optimizer.step() # 参数优化(模型训练) print('第{}轮训练,当前批次的训练损失值为:{}'.format(epoch, loss.item())) predicted = network(X_test_tensor) # 模型预测 result = predicted.data.numpy().argmax(axis=1) # 预测标签 acc_test = (result == y_test).mean() # 模型测试精度 print(acc_test) torch.save(network.state_dict(), 'leNet5-1.pt') # 保存模型参数
时间: 2024-04-22 20:28:31 浏览: 159
YOLOv5_DOTA_OBB-master-Windows运行环境配置.pdf
5星 · 资源好评率100%
这段代码是一个使用PyTorch框架实现的LeNet-5卷积神经网络模型的训练和测试过程。代码主要分为以下几个部分:
1. 数据准备部分:使用`get_data`函数获取数据,然后使用`train_test_split`将数据集划分为训练集和测试集,并将数据转换成模型要求的格式;
2. 模型定义部分:实例化一个LeNet-5网络模型;
3. 损失函数和优化器的定义:定义交叉熵损失函数和SGD优化器;
4. 模型训练部分:使用DataLoader对训练集进行批处理和打乱操作,并进行多轮迭代训练,每一轮迭代都计算一次损失值和梯度,并使用优化器更新参数;
5. 模型测试部分:使用测试集进行模型测试,并计算测试精度;
6. 模型参数保存部分:使用`torch.save`函数保存训练好的模型参数。
如果您有任何问题或需要进一步解释,请随时提出。
阅读全文