plot the loss surface (BCE and MSE)for a two layers neural network

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

def mlp_layer(x, w, b=0, activate="tanh"):
    if activate=="tanh":
        return np.tanh((w.transpose())*x + b)
    elif activate=="sigmoid":
        return 1.0/(1+np.exp(-(w.transpose()*x+b)))

def BCE(y_hat, y):
    return -(y*np.log2(y_hat)+(1-y)*np.log2(1-y_hat)).mean()

def MSE(y_hat, y):
    return ((y-y_hat)**2).mean()

def plot_face(data,data1):
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    X = np.stack([np.linspace(0, data.shape[0], data.shape[0])]* data.shape[1])
    Y = np.column_stack([np.linspace(0, data.shape[1], data.shape[1])]* data.shape[0])
    ax.plot_wireframe(X,Y,data, color="b")
    ax.plot_wireframe(X,Y,data1, color="r")
    plt.show()

x = np.random.normal(-0.5,1, 100)
x1 = np.random.normal(0.5,1, 100)
x_all= np.concatenate([x,x1], axis=0)
y_all= np.concatenate([np.zeros(100), np.ones(100)], axis=0)
out_MSE = np.zeros((100,100))
out_BCE = np.zeros((100,100))
for idx1 in range(100):
    for idx2 in range(100):
        w1 = idx1/10.0-5
        w2 = idx2/10.0-5
        out_BCE[idx1][idx2]=BCE(
            mlp_layer(w2,mlp_layer(w1,x_all), activate="sigmoid"), y_all)
        out_MSE[idx1][idx2]=MSE(
            mlp_layer(w2,mlp_layer(w1,x_all), activate="sigmoid"), y_all)

plot_face(out_MSE, out_BCE)

reference: Understanding the difficulty of training deep feedforward neural networks