srm: Can't save SRM object due to MPI

[snastase]

I want to fit brainiak.funcalign.srm.SRM on some subset of data, save this fitted SRM object to disk, then (e.g., weeks) later retrieve the fitted SRM object and use its transform_subject method to project a new subject into the shared space. I tried using numpy.save, e.g.:

from brainiak.funcalign.srm import SRM
srm = SRM(features=50)
srm.fit(data)
np.save('/path/to/saved_srm.npy', srm)

but the pickling used to save object arrays conflicts with MPI:

/jukebox/pkgs/PYGER/0.9/envs/pyger/lib/python3.6/site-packages/numpy/lib/npyio.py in save(file, arr, allow_pickle, fix_impor
ts)
    534         arr = np.asanyarray(arr)
    535         format.write_array(fid, arr, allow_pickle=allow_pickle,
--> 536                            pickle_kwargs=pickle_kwargs)
    537     finally:
    538         if own_fid:

/jukebox/pkgs/PYGER/0.9/envs/pyger/lib/python3.6/site-packages/numpy/lib/format.py in write_array(fp, array, version, allow_
pickle, pickle_kwargs)
    631         if pickle_kwargs is None:
    632             pickle_kwargs = {}
--> 633         pickle.dump(array, fp, protocol=2, **pickle_kwargs)
    634     elif array.flags.f_contiguous and not array.flags.c_contiguous:
    635         if isfileobj(fp):

TypeError: can't pickle mpi4py.MPI.Intracomm objects

Is there any other way to save the SRM object to disk and load it at a later time? If not, we should probably restructure the SRM object to make this possible. (I tried setting allow_pickle=False in numpy.save, but then it can't save the object array at all.)

[CameronTEllis]

At the moment, my solution is to save the .s_ and .w_ with numpy or python pickle calls and then remake the object. However, a better solution would be much appreciated!

[qihongl]

@snastase Thanks for opening this issue! I think making SRM object serializable would be a great enhancement.

It turns out depends on what the user wants to do, saving w and s might not be enough. For example, technically, mu_, the subject-specific intercept, should also be saved, though the current implementation forgot to use (#406). And sometimes saving rho_, sigma_ are useful as well. Being able to save the SRM object directly / avoiding doing the transformation by hand can reduce potential errors.

Currently I also just pickle all SRM parameters as np array.

[esfinn]

Have been playing around with SRM and wanted to echo that it would be great to have an easy way to save the srm object to disk for later use. Unfortunately tackling this is likely beyond my python chops at the moment, but if anyone has an idea for how to implement this, I'd be happy to help test, etc. Thanks, all, for your work on this excellent package!

[mihaic]

Fixed in #484.