Why does numpy vectorization not improve the speed of my code - python-3.x

Question

Here is the original code that does not use vectorize
import tensorflow as tf
import time
from sklearn.preprocessing import StandardScaler
from sklearn.decomposition import PCA
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data(
path='mnist.npz'
)
x_train = x_train.reshape(60000,-1)
scaler = StandardScaler()
X_scaled = scaler.fit_transform(x_train)
pca = PCA()
pca = PCA(n_components = 16) # or 12 -> 3, 4 filter_size=3
X_pca = pca.fit_transform(X_scaled).reshape(60000, 4, 4, 1)
start = time.time()
X_pca_zero = X_pca[0]
for i in range(1,60000):
X_pca_expanded = X_pca[i]
print(tf.image.ssim(X_pca_zero, X_pca_expanded, 255, filter_size=4))
print(time.time()-start)
It is essentially comparing the similarity between a reference image and a set of images. I feel it can be sped up by vectorization (so as to avoid the time wasted by the for loop). Therefore, I used the numpy vectorize function -
import numpy as np
import tensorflow as tf
import time
from sklearn.preprocessing import StandardScaler
from sklearn.decomposition import PCA
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data(
path='mnist.npz'
)
def my_func(x_train):
x_train = x_train.reshape(60000,-1)
scaler = StandardScaler()
X_scaled = scaler.fit_transform(x_train)
pca = PCA()
pca = PCA(n_components = 16) # or 12 -> 3, 4 filter_size=3
X_pca = pca.fit_transform(X_scaled).reshape(60000, 4, 4, 1)
start = time.time()
X_pca_zero = X_pca[0]
for i in range(1,6000):
X_pca_expanded = X_pca[i]
print(tf.image.ssim(X_pca_zero, X_pca_expanded, 255, filter_size=4))
print(time.time()-start)
return 0
np.vectorize(my_func(x_train))
However, there doesn't seem to be any speed improvements.

Answer 1

In Python function calls, arguments are evaluated before they are passed to the function. np.vectorize is Python function that expects a function as argument, and returns another function.
np.vectorize(my_func(x_train))
runs my_func(x_train) before passing the result to vectorize. That argument evaluation does all the prints and timing, and returns 0. I doubted if it would work, but:
In [194]: np.vectorize(0)
Out[194]: <numpy.vectorize at 0x7ff091c36310>
So it does run without error, but does nothing. All the timing is done before anything is passed to np.vectorize.
I suspect you read about the magic of "vectorization", and tried to use a like name function without reading its docs. Not only did you miss the performance disclaimer, but also didn't learn how use it (it does have its uses). It is not some sort of "compiler" or "vectorizing magic".
Most of your code is sklearn, while the only thing you time is
X_pca_zero = X_pca[0]
for i in range(1,60000):
X_pca_expanded = X_pca[i]
print(tf.image.ssim(X_pca_zero, X_pca_expanded, 255, filter_size=4))
I'm not familiar with what tf.image.ssim does. But it is a tensorflow function, and may be complex. Running it 6000 times is bound to take noticeable time. If it doesn't let you provide all X_pca values at once (as opposed to one by one), there's nothing you can do speed it up.
I don't know what it returns, but usually we don't include a print inside a timing block. Each print call adds to the computation time.