如何使用PSO优化的神经网络代码使用不同的损失函数？

Question

我正在使用pyswarms PSO进行神经网络优化。我正在尝试创建输入层和输出层的网络。

# Store the features as X and the labels as y
X = np.random.randn(25000,20)
y = np.random.random_integers(0,2,25000)


# In[29]:


def sigmoid(x):
    return 1 / (1 + math.exp(-x))


# In[58]:


print(X_train.shape)
print(y_train.shape)


# In[63]:


# Forward propagation
def forward_prop(params):
    """Forward propagation as objective function

    This computes for the forward propagation of the neural network, as
    well as the loss. It receives a set of parameters that must be 
    rolled-back into the corresponding weights and biases.

    Inputs
    ------
    params: np.ndarray
        The dimensions should include an unrolled version of the 
        weights and biases.

    Returns
    -------
    float
        The computed negative log-likelihood loss given the parameters
    """
    # Neural network architecture
    n_inputs = 20
    n_classes = 2

    # Roll-back the weights and biases
    W1 = params[0:40]

    # Perform forward propagation
    z1 = X.dot(W1)   # Pre-activation in Layer 1
    #a1 = np.tanh(z1)     # Activation in Layer 1
    #z2 = a1.dot(W2) + b2 # Pre-activation in Layer 2
    logits = z1          # Logits for Layer 2

    # Compute for the softmax of the logits
    exp_scores = np.exp(logits)
    probs = exp_scores / np.sum(exp_scores,axis=1) 

    # Compute for the negative log likelihood
    N = 25000 # Number of samples
    corect_logprobs = -np.log(probs[range(N), y])
    loss = np.sum(corect_logprobs) / N

    return loss


# In[64]:


def f(x):    
    # Compute for the negative log likelihood
    """Higher-level method to do forward_prop in the 
    whole swarm.

    Inputs
    ------
    x: numpy.ndarray of shape (n_particles, dimensions)
        The swarm that will perform the search

    Returns
    -------
    numpy.ndarray of shape (n_particles, )
        The computed loss for each particle
    """
    n_particles = x.shape[0]
    j = [forward_prop(x[i]) for i in range(n_particles)]
    return np.array(j)


# In[65]:


# Initialize swarm
options = {'c1': 0.5, 'c2': 0.3, 'w':0.9}

# Call instance of PSO
dimensions = 20   
optimizer = ps.single.GlobalBestPSO(n_particles=100, dimensions=dimensions, options=options)

# Perform optimization
cost, pos = optimizer.optimize(f, print_step=100, iters=1000, verbose=3)

我修改了示例中的代码，但是我遇到了错误

AxisError: axis 1 is out of bounds for array of dimension 1.

此外，该示例在最后一层实现softmax功能。我应该如何使用不同的损失函数？ 可以找到原始代码here。