Unable to compute multiple EMD loss

[ayushkarnawat]

Describe the bug

Attempting to use the EMD (specifically ot.emd2) to compute the loss returns a pickling error. This is likely because the loss computation could not be pickled properly (on MacOS atleast).

To Reproduce

Steps to reproduce the behavior:

1. Run code snippet below

Stack trace

Traceback (most recent call last):
  File "examples/test_emd2.py", line 25, in <module>
    emd_loss = ot.emd2(source, targets, cost)
  File "/Users/ayushkarnawat/Documents/dev/python_workspace/POT/ot/lp/__init__.py", line 429, in emd2
    res = parmap(f, [b[:, i] for i in range(nb)], processes)
  File "/Users/ayushkarnawat/Documents/dev/python_workspace/POT/ot/utils.py", line 249, in parmap
    p.start()
  File "/Users/ayushkarnawat/miniconda3/envs/pot/lib/python3.8/multiprocessing/process.py", line 121, in start
    self._popen = self._Popen(self)
  File "/Users/ayushkarnawat/miniconda3/envs/pot/lib/python3.8/multiprocessing/context.py", line 224, in _Popen
    return _default_context.get_context().Process._Popen(process_obj)
  File "/Users/ayushkarnawat/miniconda3/envs/pot/lib/python3.8/multiprocessing/context.py", line 283, in _Popen
    return Popen(process_obj)
  File "/Users/ayushkarnawat/miniconda3/envs/pot/lib/python3.8/multiprocessing/popen_spawn_posix.py", line 32, in __init__
    super().__init__(process_obj)
  File "/Users/ayushkarnawat/miniconda3/envs/pot/lib/python3.8/multiprocessing/popen_fork.py", line 19, in __init__
    self._launch(process_obj)
  File "/Users/ayushkarnawat/miniconda3/envs/pot/lib/python3.8/multiprocessing/popen_spawn_posix.py", line 47, in _launch
    reduction.dump(process_obj, fp)
  File "/Users/ayushkarnawat/miniconda3/envs/pot/lib/python3.8/multiprocessing/reduction.py", line 60, in dump
    ForkingPickler(file, protocol).dump(obj)
AttributeError: Can't pickle local object 'emd2.<locals>.f'

Screenshots

N.A.

Code sample

import numpy as np

import ot
from ot.datasets import make_1D_gauss as gauss


num_bins = 100
num_targets = 50

# Bin positions
x = np.arange(num_bins, dtype=np.float64)
means = np.linspace(20, 90, num_targets)

# Generate gaussian distributions
source = gauss(num_bins, m=20, s=5)  # m= mean, s= std
targets = np.zeros((num_bins, num_targets))
for i, m in enumerate(means):
    targets[:, i] = gauss(num_bins, m=m, s=5)

# Compute cost matrix and normalize
cost = ot.dist(x.reshape((num_bins, 1)), x.reshape((num_bins, 1)), metric="euclidean")
cost /= cost.max()

# Compute EMD loss
emd_loss = ot.emd2(source, targets, cost)

Expected behavior

The loss between each source and target is computed properly. For the example provided, this should be,

[0.0, 0.014426924576813046, 0.0288559930665439, 0.04328574159548381, 0.057715687365961325, 0.07214568550413558, 0.08657569638056341, 0.1010057100978733, 0.1154357243964259, 0.12986573880413424, 0.1442957532306665, 0.15872576766018065, 0.17315578209012886, 0.1875857965201352, 0.20201581095014873, 0.21644582538016313, 0.23087583981017767, 0.245305854240192, 0.2597358686702063, 0.2741658831002206, 0.2885958975302355, 0.30302591196024975, 0.31745592639026415, 0.3318859408202786, 0.346315955250293, 0.36074596968030737, 0.3751759841103219, 0.38960599854033634, 0.40403601297035013, 0.4184660274003606, 0.43289604183033936, 0.44732605626007876, 0.4617560706881306, 0.47618608510525495, 0.4906160994573777, 0.5050461134543907, 0.5194761256704825, 0.5339061296913891, 0.5483360991299938, 0.5627659348379075, 0.5771952970338474, 0.5916231256118063, 0.6060464162679144, 0.6204574579249451, 0.6348383926216666, 0.6491520836171457, 0.6633296042796002, 0.6772576353391107, 0.6907726671326471, 0.7036700128578843]

Environment (please complete the following information):

• OS (e.g. MacOS, Windows, Linux): MacOS
• Python version: 3.8.2
• How was POT installed (source, pip, conda): source
• Build command you used (if compiling from source): python setup.py develop
• Only for GPU related bugs:
  • CUDA version:
  • GPU models and configuration:
  • Any other relevant information:

macOS-10.13.6-x86_64-i386-64bit
Python 3.8.2 (default, May  6 2020, 02:49:43) 
[Clang 4.0.1 (tags/RELEASE_401/final)]
NumPy 1.17.4
SciPy 1.4.1
POT 0.7.0

Additional context

[ayushkarnawat]

I think the easiest way to fix this is to also check for the platform (specifically darwin) within ot/lp/__init__.py to make the number of processes equal to 1 for macOS systems (similar to windows). More specifically,

POT/ot/lp/__init__.py

Lines 375 to 376 in 793e8e1

	if sys.platform.endswith('win32'):
	processes = 1

[rflamary]

Hello,

I don't understand this. This should be covered by a test, we have a multiple process test and it runs on macosx in the tests with no errors!
https://github.com/PythonOT/POT/runs/699955225#step:6:61

In any cases, the processes=1 was a hack on windows and we should find an implementation that works on all oses instead of removing them one by one.

Since I'm only on Linux it is obviously hard for me to do that but we definitely need to be able to do that properly. Multiple EMD is an important application for people and is should work seamlessly. Maybe we should look at how sklearn do that when they allows multi-cpu computation...

[ayushkarnawat]

Why is this happening?

I think I narrowed down the issue. Both the code sample and tests work flawlessly on a local conda (or virtualenv) that uses python 3.7. This was expected since all the tests pass in the python 3.7 build you linked above. However, as soon as I switch to an env that uses python 3.8, the above code fails.

After some searching, it seems this error is linked to the changes made in the multiprocessing module for python 3.8. Specifically, on macOS (like windows), the (new) default is to start a process via a spawn rather than a fork.

This blog post explains the changes in great detail, explaining how spawning is different to forking a process. Essentially, forking copies the entire stack, which allows the parent caller to pickle the data in the variables we have access to and send that data to the child processes. On the other hand, when spawning, the child process will only inherit those resources necessary to run the process objects run() method. This means, to share resources (i.e. data) between the parent and child process, it requires that all objects/parameters passed in to be picklable. In our case, some objects are not pickable, which is why we get the error above.

Fix?

Regarding a fix, I am still unsure on how to proceed. Aaron Virshup suggests, to "run multiprocessing.set_start_method("fork") exactly once per test session." While this may be a valid workaround, the fork start method is considered unsafe as it can lead to crashes of the subprocess (see bpo-33725). This was the whole reason why the default start process was changed to spawn rather than fork in the first place.

Additionally, this approach doesn't solve the bigger problem: a function shouldn't contain objects that cannot be pickled because Windows cannot, by default, fork a process. I agree with the comment that we should find an (single) implementation that works on all OS'es. To do this, we should either (a) make all the objects/params pickable, or (b) change how the routine is run from the ground up (i.e. let the cpp code handle the multiprocessing).

[ayushkarnawat]

This also begs the question: Should we test our builds on python 3.8 for both Windows and macOS, like we do with linux already? The biggest downside is that it would increase build/CI check times.

[agramfort]

so the pb is a C++ binding that leads to code that does not pickle?

[rflamary]

It is probably the cython compiled emd solver that does not pickle indeed.

We could do the for in emd2 in parallel in the cython implementation that relaxes the GIL but I refrained from using OpenMP until now because its presence was a big question on the macosx compiler. It it better now? in this case we could also look at using the relatively new openMP emd solver from Bonneel! which would make emd faster and multicore so we would not need to do the multiple emd in parallel anymore.

[agramfort]

openmp is now supported in sklearn. See how the setup.py is done so it plays nicely with mac os

…

[ayushkarnawat]

Sorry for the late reply. I looked into the updated EMD code from Bonneel and it seems that OpenMP is required for the code to even compile, let alone run. This is different from how OpenMP is used within scipy, where the code will still work even if OpenMP is not installed within the user's system, just at the cost of speed since no parallelism is being used (obviously).

I can get started working on the port.

Edit: I also noticed that this only works with 1D distributions, should we extend this so that we can calculate EMD between higher dimensions/tensors? Or am I wrong in that the code linked also works in higher dimensions?

[rflamary]

This code works in high dimensions because we used an earlier version for POT.

[rflamary]

We have implemented multi-core EMD solver (optional) in 0.8 and we now leave the parallel implementation in the hand of the user; I'm closing this issue.