Why is this script so slow in Python?

Question

I have written a script that trains a 1D Kohonen network in MATLAB and it works as a charm. I then tried to translate it to Python 2.7, a language in which I'm pretty new, and the script is taking forever to run.

I'll explain what I'm doing to see if someone here can throw some light on the matter. I have a given dataset in the matrix y and I want to train different SOMs with it. The SOM is one-dimensional (a line), and its number of neurons varies. I train a SOM of size N=2 at first, and N=NMax at last, giving a total of NMax-2+1 SOMs. For each SOM, I want to store the weights once the training is over before moving on to the next SOM.

In MATLAB, for NMax = 5 and iterMax = 50, it takes 9.74 seconds. In Python, 54.04 seconds. This difference is huge, and the actual datasets, number of SOMs and number of iterations is even greater, so the Python code takes forever to end.

My current code is the following:

import numpy as np
import time
y = np.random.rand(2500,3) # Create random dataset to test
def A(d,s): # Neighborhood function
    return np.exp(-d**2 / (2*s**2))
sigma_0 = float(5) # Initial standard deviation for A
eta_0 = float(1) # Initial learning rate
iterMax = 250 # Maximum number of iterations
NMax = 10 # Maximum number of neurons
w = range(NMax - 1) # Initialize the size of the weight matrix (it will store NMax-2+1 sets of weights, each of varying size depending on the value of N)
#%% KOHONEN 1D
t = time.time() # Start time
for N in np.arange(2,NMax + 1): # Size of the network
    w[N - 2] = np.random.uniform(0,1,(N,np.size(y,axis=1))) - 0.5 # Initial weights
    iterCount = 1; # Iteration counter
    while iterCount < iterMax:
        # Mix the datapoints to choose them in random order
        mixInputs = y[np.random.permutation(np.size(y,axis = 0)),:]
        # Decrease the value of the variance and the learning rate
        sigma = sigma_0 - (sigma_0/(iterMax + 1)) * iterCount
        eta = eta_0 - (eta_0/(iterMax + 1)) * iterCount
        for kk in range(np.size(mixInputs,axis = 0)): # Picking one datapoint at a time
            selectedInput = mixInputs[kk,:]
            # These two lines calculate the weight that is the nearest to the datapoint selected
            aux = np.absolute(np.array(np.kron(np.ones((N,1)),selectedInput)) - np.array(w[N - 2]))
            aux = np.sum(np.abs(aux)**2,axis=-1)
            ii = np.argmin(aux) # The node ii is the winner
            for jj in range(N):
                dist = min(np.absolute(ii-jj) , np.absolute(np.absolute(ii-jj)-N)) # Centering the neighborhood function in the winner
                w[N - 2][jj,:] = w[N - 2][jj,:] + eta * A(dist,sigma) * (selectedInput - w[N - 2][jj,:]) # Updating the weights
        print(N,iterCount)
        iterCount = iterCount + 1 

elapsedTime = time.time() - t

In MATLAB, each iteration (each time the variable iterCount increases by 1) would be almost instant. In Python, each one takes a long time. I don't know why they are performing so different, but I would like to see if it is possible to speed up the Python version. Any suggestions?

EDIT: As requested in the comments, here is the much-faster MATLAB version of the code.

y = rand(2500,3) % Random dataset
A = @(d,s) exp(-d^2 / (2*s^2));
sigma_0 = 5;
eta_0 = 1;
iterMax = 250;
NMax = 10;
w = cell(NMax - 1,1);

%% KOHONEN 1D
tic();
for N = 2 : NMax 
    w{N - 1} = rand(N,size(y,2)) - 0.5;
    iterCount = 1; 
    while (iterCount < iterMax)
        mixInputs = y(randperm(size(y,1)),:);
        sigma = sigma_0 - (sigma_0/(iterMax + 1)) * iterCount;
        eta = eta_0 - (eta_0/(iterMax + 1)) * iterCount;
        for kk = 1 : size(mixInputs,1)
            input = mixInputs(kk,:);
            % [~,ii] = min(pdist2(input,w{N - 1}));
            aux = abs(repmat(input,N,1) - w{N - 1});
            [~,ii] = min((sum(aux.^2,2))); 
            for jj = 1 : N
                dist = min(abs(ii-jj) , abs(abs(ii-jj)-N));
                w{N - 1}(jj,:) = w{N - 1}(jj,:) + eta * A(dist,sigma) * (input - w{N - 1}(jj,:));
            end
        end
        N % Show N
        iterCount = iterCount + 1 % Show iterCount
    end
end
toc();

jeremycg · Accepted Answer

Here's a quick stab at some speedups - I think the output is the same, but for sure take some time to double check it:

import numpy as np
import time

np.random.seed(1234)
y = np.random.rand(2500,3) # Create random dataset to test

sigma_0 = float(5) # Initial standard deviation for A
eta_0 = float(1) # Initial learning rate
iterMax = 10 # Maximum number of iterations
NMax = 10 # Maximum number of neurons
w = {} # Initialize the size of the weight matrix (it will store NMax-2+1 sets of weights, each of varying size depending on the value of N)
#%% KOHONEN 1D
t = time.time() # Start time
for N in np.arange(2,NMax + 1): # Size of the network
    w[N - 2] = np.random.uniform(0,1,(N,np.size(y,axis=1))) - 0.5 # Initial weights
    iterCount = 1; # Iteration counter
    while iterCount < iterMax:
        # Mix the datapoints to choose them in random order
        mixInputs = y[np.random.permutation(np.size(y,axis = 0)),:]
        # Decrease the value of the variance and the learning rate
        sigma = sigma_0 - (sigma_0/(iterMax + 1)) * iterCount
        s2 = 2*sigma**2
        eta = eta_0 - (eta_0/(iterMax + 1)) * iterCount
        for kk in range(np.size(mixInputs,axis = 0)): # Picking one datapoint at a time
            selectedInput = mixInputs[kk,:]
            # These two lines calculate the weight that is the nearest to the datapoint selected
            aux = np.sum((selectedInput - np.array(w[N - 2]))**2, axis = -1)
            ii = np.argmin(aux)
            jjs = np.abs(ii - np.arange(N))
            dists = np.min(np.vstack([jjs , np.absolute(jjs-N)]), axis = 0)
            w[N - 2] = w[N - 2] + eta * np.exp((-dists**2)/s2).T[:,np.newaxis] * (selectedInput - w[N - 2]) # Updating the weights

        print(N,iterCount)
        iterCount = iterCount + 1 

elapsedTime = time.time() - t

Key speedups are mostly in using broadcasting to reduce loops/function calls.

we can replace the line:

 aux = np.absolute(np.array(np.kron(np.ones((N,1)),selectedInput)) - np.array(w[N - 2]))

with:

aux = np.abs(selectedInput - np.array(w[N - 2]))

(I have further crammed this in to a couple of the next steps). Numpy broadcasting gives us the same result, without have to take the kron product.

As an example:

np.kron(np.ones((3,1)), np.array([6,5,4])) - np.arange(-9,0).reshape(3,3)

is the same output as:

np.array([6,5,4]) - np.arange(-9,0).reshape(3,3)

kron(np.ones(N,1),x) gives an N * x.shape[0] array with N copies of x. Broadcasting takes care of this in a much cheaper manner.

The other main speed up is to reduce the:

for jj in range(N):

To a matrix operation. We calculate the 2*sigma**2 once per loop, replace the A function with a native numpy call, and vectorize the rest.

Why is this script so slow in Python?

Answers (2)

Related Questions