Optimizing Numpy Operations - python-3.x

Question

I am trying to train a multi-class classifier with multinomial logistic regression and gradient descent. Specifically, the model will have a trained weights matrix w with shape (C, D) where C is the number of classes and D is the number of features of each input. Also, we will have a bias vector b with dimension (C,). We have an (N, D) input matrix X, where N is the number of training inputs, and a vector y with shape (N,), where each entry in y is a number from 0 to C - 1, indicating which class the input belongs to. I have written the following code:
for _ in range(max_iterations):
z = np.apply_along_axis(lambda v: v - max(v), 1, X # w.T + b)
probs = np.exp(z)
denom = np.sum(probs, axis=1)
for i in range(C):
for j in range(N):
if i == y[j]:
w[i] -= (step_size / N) * ((probs[j][i] / denom[j]) - 1) * X[j]
b[i] -= (step_size / N) * ((probs[j][i] / denom[j]) - 1)
else:
w[i] -= (step_size / N) * (probs[j][i] / denom[j]) * X[j]
b[i] -= (step_size / N) * (probs[j][i] / denom[j])
This produces the correct weights and bias that I want, but clearly it doesn't take advantage of numpy's operations to speed things up. So I tried to speed some of it up with the following code:
for _ in range(max_iterations):
z = np.apply_along_axis(lambda v: v - max(v), 1, X # w.T + b)
probs = np.exp(z)
denom = np.sum(probs, axis=1)
s = np.zeros((N, C))
for i in range(N):
s[i] = probs[i] / denom[i]
for i in range(N):
s[i][y[i]] += -1
for c in range(C):
grad_w = s.T[c] # X
w[c] += (step_size / N) * grad_w
b[c] += (step_size / N) * sum(s.T[c])
I was hoping that this would produce the same results as in the previous part while being faster... and it managed to be faster, but with incorrect results.
So I have a couple of questions. First, why is my second piece of code not producing the right results, and what would be a fix for it? Second, and more importantly, how would I optimize this further? This is mainly for me to learn how to take advantage of numpy's vectorized operations.

Answer 1

This may help with some of the iterations.
Start with a small 2d array:
In [251]: probs = np.arange(12).reshape(3,4)
In [252]: denom = np.sum(probs, axis=1)
In [253]: denom
Out[253]: array([ 6, 22, 38])
To divide a (3,4) array by a (3,), we need to make the later (3,1):
In [254]: probs/denom[:,None]
Out[254]:
array([[0. , 0.16666667, 0.33333333, 0.5 ],
[0.18181818, 0.22727273, 0.27272727, 0.31818182],
[0.21052632, 0.23684211, 0.26315789, 0.28947368]])
Read, and reread, the numpy documentation on broadcasting if that doesn't make sense.
Another way to get the required 2d denom, is:
In [255]: denom = np.sum(probs, axis=1, keepdims=True)
In [256]: denom
Out[256]:
array([[ 6],
[22],
[38]])
In [257]: probs/denom
Out[257]:
array([[0. , 0.16666667, 0.33333333, 0.5 ],
[0.18181818, 0.22727273, 0.27272727, 0.31818182],
[0.21052632, 0.23684211, 0.26315789, 0.28947368]])
The same should work for the max subtraction that you use with apply_along_axis. apply... is not a speed tool, and not superior to simple iteration.
In [258]: np.max(probs, axis=1, keepdims=True)
Out[258]:
array([[ 3],
[ 7],
[11]])
In [259]: probs - _
Out[259]:
array([[-3, -2, -1, 0],
[-3, -2, -1, 0],
[-3, -2, -1, 0]])