05 Pytorch代码实践之【多元分类】

import torch
import matplotlib.pyplot as plt
from torch import nn, optim
from time import perf_counter

cluster = torch.ones(500, 2)
data0 = torch.normal(4 * cluster, 2)
data1 = torch.normal(-4 * cluster, 1)
data2 = torch.normal(-8 * cluster, 1)
label0 = torch.zeros(500)
label1 = torch.ones(500)
label2 = label1 * 2
x = torch.cat((data0, data1, data2)).type(torch.FloatTensor)
y = torch.cat((label0, label1, label2)).type((torch.LongTensor))

plt.scatter(x.numpy()[:, 0], x.numpy()[:, 1], c=y.numpy(), s=10, lw=0, cmap='RdYlGn')
plt.show()


class Net(nn.Module):
    def __init__(self, input_feature, num_hidden, outputs):
        super(Net, self).__init__()
        self.hidden = nn.Linear(input_feature, num_hidden)
        self.out = nn.Linear(num_hidden, outputs)

    def forward(self, x):
        x = torch.nn.functional.relu(self.hidden(x))
        x = self.out(x)
        x = torch.nn.functional.softmax(x)
        return x


CUDA = torch.cuda.is_available()
if CUDA:
    net = Net(input_feature=2, num_hidden=20, outputs=3).cuda()
    inputs = x.cuda()
    target = y.cuda()
else:
    net = Net(input_feature=2, num_hidden=20, outputs=3)
    inputs = x
    target = y

optimizer = optim.SGD(net.parameters(), lr=0.02)
criterion = nn.CrossEntropyLoss()


def draw(output):
    if CUDA:
        output = output.cpu()  # 还原为cpu类型才能进行绘图
    plt.cla()  # 清空画布
    output = torch.max(output, 1)[1]  # [max_value,max_value_index][1]->index 最大值索引（0或1）
    pre_y = output.data.numpy().squeeze()  # pre_y = output.numpy() 也可以的
    target_y = y.numpy()
    plt.scatter(x.numpy()[:, 0], x.numpy()[:, 1], c=pre_y, s=10, lw=0, cmap='RdYlGn')
    accuracy = sum(pre_y == target_y) / 1500.0
    plt.text(1.5, -4, 'Accuracy=%s' % (accuracy),
             fontdict={'size': 20, 'color': 'red'})
    plt.pause(0.1)


def train(model, criterion, optimizer, epochs):
    for epoch in range(epochs):
        output = model(inputs)

        loss = criterion(output, target)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if epoch % 40 == 0:
            draw(output)
    return model, loss


start = perf_counter()
model, loss = train(net, criterion, optimizer, epochs=5000)
finish = perf_counter()
time = finish - start
print("计算时间：%s" % time)
print("final loss:", loss.item())
print("weights:", list(model.parameters()))

结果：

…

…

计算时间：41.4019583
final loss: 0.5656905770301819
weights: [Parameter containing:
tensor([[-0.3741, -0.7318],
        [ 0.4266,  0.5075],
        [ 0.7998,  0.6104],
        [-0.3695, -0.0914],
        [-0.3085, -0.5864],
        [ 0.1329,  0.3026],
        [ 0.7320,  0.2708],
        [-0.7223,  0.0390],
        [-0.1677, -0.8182],
        [-0.3276, -0.3438],
        [ 0.4860, -0.5646],
        [-0.4681, -0.3976],
        [-0.0435,  0.1630],
        [ 0.6001,  0.7026],
        [ 0.6207,  0.5880],
        [ 0.2340, -0.7142],
        [ 0.4925, -0.1610],
        [-0.2655,  0.2880],
        [ 0.5097,  0.4200],
        [-0.4757, -0.1096]], device='cuda:0', requires_grad=True), Parameter containing:
tensor([ 0.5086,  0.6349, -0.5179, -0.8939,  1.0239,  0.0568,  0.4589, -1.0071,
        -1.6049,  0.6958,  0.9632, -0.7235, -0.4089, -0.0101,  0.1009, -0.7955,
         0.0696,  1.0634,  0.0237,  0.5647], device='cuda:0',
       requires_grad=True), Parameter containing:
tensor([[-0.4928,  0.2230,  0.2948, -0.2385, -0.2292,  0.0068,  0.4368, -0.2500,
         -0.0082, -0.1561,  0.2489,  0.0018,  0.1111,  0.4554,  0.3450, -0.0349,
         -0.0302,  0.2761,  0.1810, -0.2328],
        [ 0.2713, -0.1252, -0.4039, -0.5439,  0.6500,  0.0083, -0.2000, -0.6431,
         -1.1661,  0.6345,  0.5225, -0.3631,  0.0581, -0.1914, -0.1908, -0.1370,
         -0.2039,  0.6197, -0.3525,  0.2870],
        [-0.1955, -0.1500, -0.0243,  0.6625, -0.4319, -0.2280, -0.2149,  0.7826,
          1.1261, -0.3013, -0.5921,  0.6576,  0.1273,  0.0788, -0.3057,  0.5231,
          0.0765, -0.6810,  0.0587, -0.3855]], device='cuda:0',
       requires_grad=True), Parameter containing:
tensor([ 0.3078,  0.7959, -1.2601], device='cuda:0', requires_grad=True)]