整体代码# coding:utf8# 解决 OpenMP 库冲突问题importos os.environ[KMP_DUPLICATE_LIB_OK]TRUEimporttorchimporttorch.nnasnnimportnumpyasnpimportrandomimportjsonimportmatplotlib.pyplotasplt 基于pytorch框架编写模型训练 实现一个自行构造的找规律(机器学习)任务 规律x是一个5维向量如果第1个数第5个数则为正样本反之为负样本 classTorchModel(nn.Module):def__init__(self,input_size):super(TorchModel,self).__init__()self.linearnn.Linear(input_size,1)# 线性层self.activationtorch.sigmoid# nn.Sigmoid() sigmoid归一化函数self.lossnn.functional.mse_loss# loss函数采用均方差损失# 当输入真实标签返回loss值无真实标签返回预测值defforward(self,x,yNone):xself.linear(x)# (batch_size, input_size) - (batch_size, 1)y_predself.activation(x)# (batch_size, 1) - (batch_size, 1)ifyisnotNone:returnself.loss(y_pred,y)# 预测值和真实值计算损失else:returny_pred# 输出预测结果# 生成一个样本, 样本的生成方法代表了我们要学习的规律# 随机生成一个5维向量如果第一个值大于第五个值认为是正样本反之为负样本defbuild_sample():xnp.random.random(5)ifx[0]x[4]:returnx,1else:returnx,0# 随机生成一批样本# 正负样本均匀生成defbuild_dataset(total_sample_num):X[]Y[]foriinrange(total_sample_num):x,ybuild_sample()X.append(x)Y.append([y])# print(X)# print(Y)returntorch.FloatTensor(X),torch.FloatTensor(Y)# 测试代码# 用来测试每轮模型的准确率defevaluate(model):model.eval()test_sample_num100x,ybuild_dataset(test_sample_num)print(本次预测集中共有%d个正样本%d个负样本%(sum(y),test_sample_num-sum(y)))correct,wrong0,0withtorch.no_grad():y_predmodel(x)# 模型预测 model.forward(x)fory_p,y_tinzip(y_pred,y):# 与真实标签进行对比iffloat(y_p)0.5andint(y_t)0:correct1# 负样本判断正确eliffloat(y_p)0.5andint(y_t)1:correct1# 正样本判断正确else:wrong1print(正确预测个数%d, 正确率%f%(correct,correct/(correctwrong)))returncorrect/(correctwrong)defmain():# 配置参数epoch_num20# 训练轮数batch_size20# 每次训练样本个数train_sample500# 每轮训练总共训练的样本总数input_size5# 输入向量维度learning_rate0.01# 学习率# 建立模型modelTorchModel(input_size)# 选择优化器optimtorch.optim.Adam(model.parameters(),lrlearning_rate)log[]# 创建训练集正常任务是读取训练集train_x,train_ybuild_dataset(train_sample)# 训练过程forepochinrange(epoch_num):model.train()watch_loss[]forbatch_indexinrange(train_sample//batch_size):#取出一个batch数据作为输入 train_x[0:20] train_y[0:20] train_x[20:40] train_y[20:40]xtrain_x[batch_index*batch_size:(batch_index1)*batch_size]ytrain_y[batch_index*batch_size:(batch_index1)*batch_size]lossmodel(x,y)# 计算loss model.forward(x,y)loss.backward()# 计算梯度optim.step()# 更新权重optim.zero_grad()# 梯度归零watch_loss.append(loss.item())print(\n第%d轮平均loss:%f%(epoch1,np.mean(watch_loss)))accevaluate(model)# 测试本轮模型结果log.append([acc,float(np.mean(watch_loss))])# 保存模型torch.save(model.state_dict(),model.bin)# 画图print(log)plt.plot(range(len(log)),[l[0]forlinlog],labelacc)# 画acc曲线plt.plot(range(len(log)),[l[1]forlinlog],labelloss)# 画loss曲线plt.legend()plt.show()return# 使用训练好的模型做预测defpredict(model_path,input_vec):input_size5modelTorchModel(input_size)model.load_state_dict(torch.load(model_path))# 加载训练好的权重print(model.state_dict())model.eval()# 测试模式withtorch.no_grad():# 不计算梯度resultmodel.forward(torch.FloatTensor(input_vec))# 模型预测forvec,resinzip(input_vec,result):print(输入%s, 预测类别%d, 概率值%f%(vec,round(float(res)),res))# 打印结果if__name____main__:# main()test_vec[[0.81889086,0.15229675,0.31082123,0.03504317,0.88920843],[0.94963533,0.5524256,0.95758807,0.95520434,0.84890681],[0.90797868,0.67482528,0.13625847,0.34675372,0.19871392],[0.39349776,0.59416669,0.92579291,0.41567412,0.8358894]]predict(model.bin,test_vec)具体来说它实现了以下核心功能任务定义这是一个“找规律”的分类任务。输入一个 5 维的随机向量xxx。规律如果向量的第 1 个数索引 0大于第 5 个数索引 4则标记为正样本标签为 1否则为负样本标签为 0。模型构建 (TorchModel类)使用了一个简单的线性层 (nn.Linear) 将 5 维输入映射到 1 维输出。使用sigmoid函数将输出归一化到 0-1 之间用于二分类概率预测。损失函数采用均方差损失 (MSELoss)。数据生成build_sample函数随机生成符合上述规律的单个样本。build_dataset函数批量生成正负样本均匀的训练数据。训练流程 (main函数)配置了训练轮数 (epoch)、批次大小 (batch_size)、学习率等超参数。使用 Adam 优化器进行模型训练。每轮训练后评估模型在测试集上的准确率并记录 Loss 和 Accuracy 变化最后绘制曲线图并保存模型权重 (model.bin)。预测功能 (predict函数)加载训练好的模型对给定的 5 维向量输入进行分类预测并输出预测类别及概率值。代码底部包含了一段测试代码用于演示如何加载模型并对 4 组示例数据进行预测。