Using dask with PyTorch (train a model)

I find Dask very useful, and would like to use it to accelerate some pytorch models I am working on. A function to be accelerated would be the forward() function. I started doing some tests, but I get errors like these:

  File "/home/muammar/.local/lib/python3.7/site-packages/distributed/protocol/serialize.py", line 81, in serialization_error_loads
    raise TypeError(msg)
TypeError: Could not serialize object of type Tensor.
Traceback (most recent call last):
  File "/home/muammar/.local/lib/python3.7/site-packages/distributed/protocol/serialize.py", line 144, in serialize
    header, frames = dumps(x, context=context) if wants_context else dumps(x)
  File "/home/muammar/.local/lib/python3.7/site-packages/distributed/protocol/serialize.py", line 38, in dask_dumps
    header, frames = dumps(x)
  File "/home/muammar/.local/lib/python3.7/site-packages/distributed/protocol/torch.py", line 10, in serialize_torch_Tensor
    header, frames = serialize(t.numpy())
RuntimeError: Can't call numpy() on Variable that requires grad. Use var.detach().numpy() instead.

If I detach the variable, then pytorch does not work with numpy arrays. My question would be – what is the state of pytorch support? If it is of any help, I would be more than happy to try any code (or even write something if it is under my capabilities) because it would be very useful to have this support in Dask. Thanks for this great python module.

[mrocklin]

Thank you for thinking of us @muammar . Are you able to share the article?

[stsievert]

I would like to add you in the acknowledgment section.

@muammar no objections here! Thank you.