How to use Dask to parallelize QNode computations?

Hi everyone. I’m working on quantum neural networks for computer vision. I wonder if there is a better or more efficient way to implement quantum version of convolution rather than using a number of for loops like the one demonstrated in this tutorial? Josh suggested me to use Dask to parallelize QNode computations (Thanks @josh again! ). However, it seems I did not use it correctly because when I executed the demo codes provided in the documentation of Dask, I found the computation got even much slower. Here is the code I run:

import dask
import numpy as np
def inc(x):
    return x + 1
def double(x):
    return x * 2
def add(x, y):
    return x + y
data = np.random.rand((100))
t0 = time.time()
output = []
for x in data:
    a = dask.delayed(inc)(x)
    b = dask.delayed(double)(x)
    c = dask.delayed(add)(a, b)
    output.append(c)
total = dask.delayed(sum)(output)
total.compute() 
t1 = time.time()
print("Time: ", (t1 - t0) )
t0 = time.time()
output = []
for x in data:
    a = (inc)(x)
    b = (double)(x)
    c = (add)(a, b)
    output.append(c)
total = (sum)(output)
t1 = time.time()
print("Time: ", (t1 - t0) )
Time:  0.034467220306396484
Time:  0.0003325939178466797

Could anyone help me with it? Many thanks!

Hi @gojiita_ku,

Welcome! :slightly_smiling_face:

The QNodeCollection objects could be helpful in parallelizing QNode evaluations. We have Dask support for a QNodeCollection by specifying the parallel=True option (see the Asynchronous evaluation in the related docs).

It is worth noting, however, that as the documentation mentions, this option will only be useful with some devices (e.g., can be beneficial with the QVM device from the Forest plugin, as used in this tutorial).

Let us know how this goes! :slightly_smiling_face:

Hi @antalszava,

Thanks for your answer! I’ll try QNodeCollection for my case and see if it would work.:grinning:

I’m very excited to see Pennylane has been updated so quickly! I’m currently only using Pennylane for QML.

1 Like

Hi @antalszava,

I checked the documentation of QNodeCollection. It seems like all QNodes within a QNodeCollection must have the same input. But if we use QNodeCollection in the case of quantum convolution, all QNodes would have different inputs (e.g. different patches in the two-dimensional image). So that means QNodeCollection would not work for this case? I’m not sure if I got a correct interpretation of the class QNodeCollection.

Hi @gojiita_ku,

Indeed, that’s right. A QNodeCollection assumes that the inputs are the same for all the QNodes.

  1. Could potentially all QNode inputs be concatenated into a single object, which is than passed to the QNodeCollection? Indexing into the input could then be a way for distributing the input parameters. Something along the lines of
import pennylane as qml
import numpy as np

dev = qml.device('default.qubit', wires=2)

@qml.qnode(dev)
def circ1(par):
    qml.RY(par[0], wires=0)  # <---- 1. parameter
    return qml.expval(qml.PauliZ(0))

@qml.qnode(dev)
def circ2(par):
    qml.RX(par[1], wires=0)  # <---- 2. parameter
    return qml.expval(qml.PauliZ(0))

qnodes = qml.QNodeCollection([circ1, circ2])

par1 = np.array([0.1234])
par2 = np.array([0.4323])

pars = np.concatenate([par1, par2]) # <---- Concatenating input parameters

qnodes(pars)

There is a historic reason for why a QNodeCollection behaves like this, mainly circuits with similar structures were considered when creating it.

  1. The Dask logic in QNodeCollection is as follows:
for q in self.qnodes:
    results.append(dask.delayed(q)(*args, **kwargs))

return dask.compute(*results, scheduler=_scheduler)

Here, _scheduler is a Dask scheduler to use, can be for example "threads".

Alternatively, this could also help with a custom solution to using Dask with multiple qnodes. The input parameters for each qnode could then be zipped with the qnode itself (just an idea for a potential approach):

for args, q in zip(arg_lists, qnodes):
    results.append(dask.delayed(q)(*args, **kwargs))

where arg_lists would be an ordered list of arguments for each qnode and qnodes are the QNodes to evaluate.

Hope some of this is helpful, let us know how it goes! :slightly_smiling_face:

Hi @antalszava,

Thank you so much! I’ll try both of your suggestions and get back to you soon. :grinning:

Sounds good! :slightly_smiling_face:

1 Like