Visualization of two-dimensional feature logistic regression prediction results

Problem description
The following 2D features , Each sample is a positive sample （ Red ） Or negative sample （ Blue ）, Implement the binary classification model

Single layer neural network
$$\left[\begin{array}{cc}{w}_1& w_2\end{array}\right]\left[\begin{array}{c}{x}_1\\ x_2\end{array}\right]+\left[\begin{array}{c}b\end{array}\right]=\left[\begin{array}{c}y\end{array}\right]$$

class LogisticRegression(nn.Module):
    def __init__(self):
        super(LogisticRegression, self).__init__()
        self.lr = nn.Linear(2, 1)
        self.sm = nn.Sigmoid()

    def forward(self, x):
        out = self.lr(x)
        out = self.sm(out)
        return out

We want to see the training effect of the model , Need to put $Linear$ Layer Visualization
$$w_1\ast x_1+w_2\ast x_2+b=y$$
$y=0$ when ,$x_1$ and $x_2$ The relationship is as follows ：
$$x_2=\frac{-w_1\ast x_1-b}{w2}$$
take $x_1$ stay $[30,100]$ Between , Using neural networks $Linear$ Layer parameters will $x_2$ Find the value of , Draw the training results .

def vis_one_layer(logistic_model):
    w1, w2 = logistic_model.lr.weight.data.numpy()[0]
    b = logistic_model.lr.bias.data.numpy()[0]
    plot_x = np.arange(30, 100, 0.1)
    plot_y = (-w1 * plot_x - b) / w2
    plt.plot(plot_x, plot_y)
    plt.show()

Multilayer neural network
If the neural network has more than one layer , It is impossible to substitute $x_1$ and $x_2$ The relationship between , How to visualize the prediction results of neural networks ？

class MyLogistic(nn.Module):
    def __init__(self, input_size):
        super().__init__()
        self.hidden_1 = nn.Linear(input_size, 64)
        self.hidden_2 = nn.Linear(64, 32)
        self.hidden_3 = nn.Linear(32, 16)
        self.output = nn.Linear(16, 1)
        self.relu_1 = nn.ReLU()
        self.relu_2 = nn.ReLU()
        self.relu_3 = nn.ReLU()
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        x = self.hidden_1(x)
        x = self.relu_1(x)
        x = self.hidden_2(x)
        x = self.relu_2(x)
        x = self.hidden_3(x)
        x = self.relu_3(x)
        x = self.output(x)
        x = self.sigmoid(x)
        return x

The idea is to use a series of points to determine the region of the neural network to divide the binary classification problem , Just because the predicted result is not linear , So use the model to predict the positive and negative for the points in the region , Then visualize with different colors , Use $plt.contourf$ Draw outline lines and fill .

def vis_result(x_data, y_data, model):
    x_min, x_max = x_data[:, 0].min() - 1, x_data[:, 0].max() + 1
    y_min, y_max = x_data[:, 1].min() - 1, x_data[:, 1].max() + 1
    print(x_min, x_max, y_min, y_max)
    h = 0.1
    xx, yy = np.meshgrid(np.arange(x_min, x_max, h), np.arange(y_min, y_max, h))
    z = model(torch.from_numpy(np.c_[xx.ravel(), yy.ravel()]).float())
    z = z.reshape(xx.shape).detach().numpy()
    print(z)
    z = [z[i] >= 0.5 for i in range(len(z))]
    z = np.array(z)
    plt.contourf(xx, yy, z, alpha=0.3)
    for i in range(len(y_data)):
        if y_data[i] == 1:
            plt.scatter(x_data[i][0], x_data[i][1], c='r')
        else:
            plt.scatter(x_data[i][0], x_data[i][1], c='b')
    plt.show()