Dataloader pickle error because *must* use multiprocessing_context = 'spawn' instead of 'fork'

data
1

I’m working with a library that does some heavy multithreading and it hangs if I use the ‘fork’ multiprocessing context, so I need to use ‘spawn’ (not using windows jfc…).

Then getting the classic “AttributeError: Can’t pickle local object ‘main..collate_fn’” when defining the collate function something like this inside the main():

def collate_fn(x):
    return functions.collate(collate_params, x)

What can I do?

Someone somewhere mentioned using dill instead of pickle (it can serialize anything I guess), but this is a problem with pytorch using pickle, so can’t just use dill I think…

EDIT: OK the form of the collate function is not really important, but that it’s defined inside the main() and then right after I define the Dataloader(collate_fn=collate_fn, ...)

2

Can you elaborate on what is making the function unpicklable?

One potentially solution is to create a callable object with custom __getstate__ and __setstate__ methods that will make that the object pickable (or use dill within). Then use that callable object as collate_fn.

1 Like
3

Ah OK that sounds cool, need to try that out!

I don’t know why it’s unpicklable… it’s not a lambda expression after all.

Anyway, I actually managed to use ‘fork’ by going around the multiprocessing issue so AOK (will stick with ‘fork’ if I only can).

Thanks for the interesting tip though @nivek!