I'm trying to refactor some code for fitting parametric models that are defined symbolically using Theano. My goal is for all models to expose a common interface so that, as far as possible, they can be drop-in replacements for one another. To that end I've tried to encapsulate each model within a separate class. A key requirement is that I'm able to parallelize fitting the same model to multiple datasets using multiprocessing (I'm using the joblib wrapper for this).Here's a runnable example of what I'm doing at the moment:import numpy as npimport theanofrom theano import tensor as tefrom theano.gradient import jacobian, hessianfrom scipy.optimize import minimizefrom joblib import Parallel, delayedclass Rosenbrock(object):    The Rosenbrock function: f(x, y) = (a - x)^2 + b(y - x^2)^2     # symbolic variables - only used internally    _P = te.dvector('P')    _a, _b = _P[0], _P[1]    _xy = te.dmatrix('xy')    _x, _y = _xy[0], _xy[1]    _z = te.dvector('z')    _z_hat = (_a - _x) ** 2 + _b * (_y - _x ** 2) ** 2    _diff = _z - _z_hat    _loss = 0.5 * te.dot(_diff, _diff)    _jac = jacobian(_loss, _P)    _hess = hessian(_loss, _P)    # theano functions - part of the interface    forward = theano.function([_P, _xy], _z_hat)    loss = theano.function([_P, _xy, _z], _loss)    jacobian = theano.function([_P, _xy, _z], _jac)    hessian = theano.function([_P, _xy, _z], _hess)    @staticmethod    def initialize(xy, z):                make some sensible estimate of what the initial parameters should be,        based on xy and z                P0 = xy[:, np.argmin(z)]        return P0    @staticmethod    def _postfit(P):                sometimes I want to make some adjustments to the parameters post-        fitting, e.g. wrapping angles between 0 and 2pi                return Pdef do_fit(model, *args):        wrapper function that performs the fitting        # initialize the model    P0 = model.initialize(*args)    # do the fit    res = minimize(model.loss, P0, args=args, method='Newton-CG',                   jac=model.jacobian, hess=model.hessian)    P = res.x    # tweak the parameters    P = model._postfit(P)    # return the tweaked parameters    return Pdef run(niter=2000):    # I don't actually need to instantiate this, since everything is    # effectively a class method...    model = Rosenbrock()    # some example data    xy = np.mgrid[-3:3:100j, -3:3:100j].reshape(2, -1)    P = np.r_[1., 100.]    z = model.forward(P, xy)    # run multiple fits in parallel    pool = Parallel(n_jobs=-1, verbose=1, pre_dispatch='all')    results = pool(delayed(do_fit)(model, xy, z) for _ in xrange(niter))if __name__ == '__main__':    run()The core functions forward(),  loss(), jacobian() and hessian() behave like staticmethods of the class. I've found that in order to be able to parallelizethe fitting, the Theano functions must be attributes of the class rather thanof an instance. Otherwise (i.e. if I define these functions within the __init__()method of the class), what happens when I try to call them in parallel usingmultiprocessing is that the worker threads block one another, meaning that they effectively only use a single core. Presumably this is because the GIL is no longer being circumvented, although I don't really understand why this should be the case.Declaring the Theano functions as class methods has two rather undesirable consequences:The class namespace gets cluttered up with all of the intermediateTheano symbolic variables (a, b etc.), which aren't really neededafter theano.function() has been called. To hide them from the user I haveto prepend the variable names with underscores, which makes the code harderto read.theano.function() triggers the auto-generation and compilation ofC code, which is a slow process. It would be most convenient to defineall of my models in the same source file, but that means that wheneverI import or reload that file I have to wait for all of my models tore-compile.Can anyone (especially someone with experience in Theano) suggest a better wayto structure this code?
Parallelizing some model fitting code using multiprocessing
python;object oriented;multithreading;numpy