鸢尾花分类-pytorch实现

前言

本文用pytorch实现了鸢尾花分类，数据不多，只做代码展示用，后续有升级版本。

代码

'''
-*- coding: utf-8 -*-
@File  : main.py
@Author: Shanmh
@Time  : 2024/05/06 上午9:37
@Function：
'''
import torch
from sklearn import datasets
import torch.nn as nn

#1.数据准备
dataset=datasets.load_iris()
print(dataset["data"][:10])
print(dataset["target"][:10])
i_data=torch.FloatTensor(dataset["data"])
i_target=torch.LongTensor(dataset["target"])

#2.模型构建
class IrisModel(nn.Module):
    def __init__(self,input_n=4,hidden_n=20,output_n=3):
        super().__init__()
        self.line1=nn.Linear(input_n,hidden_n)
        self.line2=nn.Linear(hidden_n,output_n)
        self.relu=nn.ReLU()
    def forward(self,x):
        x=self.line1(x)
        x=self.relu(x)
        x=self.line2(x)
        return x

#3.参数定义
epoch=500
lr=0.01

model=IrisModel()
optimizer=torch.optim.SGD(model.parameters(),lr=lr) #定义优化器
loss_fun=torch.nn.CrossEntropyLoss() #多分类采用交叉熵损失函数


for e in range(epoch):
    out=model(i_data)
    loss=loss_fun(out,i_target)
    optimizer.zero_grad()  # 梯度清零
    loss.backward()  # 前馈操作
    optimizer.step()


# 5. 得出结果
out = model(i_data)
prediction = torch.max(out, 1)[1]
pred_y = prediction.data.numpy()
target_y = i_target.data.numpy()
result=pred_y==target_y
print(f"模型预测准确度，acc：{'{:.2f}'.format(len(result[result==True])/len(result))}%")