Iris Classification on PyTorch

code

# -*- coding:utf8 -*-from sklearn.datasets import load_irisfrom sklearn.utils import shufflefrom sklearn.model_selection import train_test_splitimport torchimport torch.optim as optimimport torch.nn as nnimport torch.nn.functional as Fhl = 10lr = 0.005num_epoch = 50000class Net(nn.Module):    def __init__(self):        super(Net, self).__init__()        self.fc1 = nn.Linear(4, hl)        self.fc2 = nn.Linear(hl, 3)        self.softmax = torch.nn.Softmax(dim=1)    def forward(self, x):        out = self.fc1(x)        out = F.relu(out)        out = self.fc2(out)        out = self.softmax(out)        return outif __name__ == '__main__':    iris = load_iris()    x, y = shuffle(iris.data,iris.target)    x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.25, random_state=33)    net = Net()    criterion = nn.CrossEntropyLoss()    optimizer = optim.Adam(net.parameters(), lr=lr)    # train    for epoch in range(num_epoch):        x = torch.Tensor(x_train).float()        y = torch.Tensor(y_train).long()        optimizer.zero_grad()        y_pred = net(x)        loss = criterion(y_pred, y)        loss.backward()        optimizer.step()        if epoch % 50 is 0:            print(loss) # cross entropy    # test    x = torch.Tensor(x_test).float()    y = torch.Tensor(y_test).long()    y_pred = net(x)    _, predicted = torch.max(y_pred, 1)    acc = torch.sum(y == predicted).numpy() / len(x_test)    print(acc)