pytorch做一个分类任务

来两个圈吧

import matplotlib.pyplot as plt
from sklearn.datasets import make_circles

n_samples = 800
X, y = make_circles(n_samples=n_samples, noise=0.1, random_state=1, factor=0.6)
plt.scatter(X[:, 0], X[:, 1], c=y)

建立神经网络

import torch
import torch.nn.functional as F

class Net(torch.nn.Module):
    def __init__(self, n_feature, n_hidden, n_hidden2, n_output):
        super().__init__()
        self.hidden = torch.nn.Linear(n_feature, n_hidden)
        self.hidden2 = torch.nn.Linear(n_hidden, n_hidden2)
        self.out = torch.nn.Linear(n_hidden2, n_output)

    def forward(self, x):
        x = F.relu(self.hidden(x))
        x = torch.sigmoid(self.hidden2(x))
        x = self.out(x)
        x = F.softmax(x, dim=1)
        return x