【实验】鸢尾花分类——简单的神经网络

import torch
from torch import nn
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
import numpy as np
import matplotlib.pyplot as plt

X = torch.tensor(load_iris().data, dtype=torch.float32)  
y = torch.tensor(load_iris().target, dtype=torch.long)

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)


导入鸢尾花数据集，这里注意数据和标签类型的设置：dtype=torch.float32，dtype=torch.long，否则会报错

net = nn.Sequential(nn.Linear(4, 10), 
　　　　　　　　　　　　nn.ReLU(),
　　　　　　　　　　　　nn.Linear(10, 10),
　　　　　　　　　　　　nn.ReLU()，
　　　　　　　　　　　　nn.Linear(10, 3))

def init_weights(m):
    if type(m) == nn.Linear:
        nn.init.normal_(m.weights, std=0.01)

loss = nn.CrossEntropyLoss(reduction="none")      

trainer = torch.optim.Adam(net.parameters(), lr=0.05)

train_loss = []
test_loss = []
train_l = loss(net(X_train), y_train).mean().detach().numpy()
test_l = loss(net(X_test), y_test).mean().detach().numpy()
train_loss.append(train_l)
test_loss.append(test_l)

epochs = 1000

for i in range(epochs):
    trainer.zero_grad()
    l = loss(net(X_train), y_train).mean()
    l.backward()
    trainer.step()
    l = loss(net(X), y).mean()
    
    train_l = loss(net(X_train), y_train).mean().detach().numpy()
    test_l = loss(net(X_test), y_test).mean().detach().numpy()
    train_loss.append(train_l)
    test_loss.append(test_l)
epoch_index = range(epochs + 1)
plt.plot(epoch_index, train_loss, 'green',  epoch_index, test_loss, 'blue')   
plt.show()