QNGOptimizer for PyTorch

Hi , i am trying to use the QNGOptimizer as a custom optimizer for pytorch. I followed the example given in (Quantum Natural Gradient Descent - #2 by josh) but i am using a torchlayer which is making it difficult somehow for it to find the qnode.

In my closure() i return the metric_tensor as:
return loss, qml.metric_tensor(policy_estimator.qlayer)

and i have defined the QNGOptimizer as :

class QNGOptimizer(torch.optim.Optimizer):

    def __init__(self, params, lr=0.01, diag_approx=False, lam=0):
        defaults = dict(lr=0.01, diag_approx=False, lam=0)
        super().__init__(params, defaults)

    def step(self, closure=None):
        loss = None

        if closure is not None:
            loss, metric_tensor = closure()

        for group in self.param_groups:
            for p in group["params"]:

                if p.grad is None:

                grad = p.grad.data
                state = self.state[p]

                # State initialization
                if len(state) == 0:
                    state["step"] = 0

                g = metric_tensor(p)
                g += group["lam"] * np.identity(g.shape[0])

                state["step"] += 1

                d_p = torch.tensor(-group['lr'] * np.linalg.solve(g, grad))

        return loss

I am having an error starting in the line g = metric_tensor(p) with the following message:

g = metric_tensor(p)# approx=group[“diag_approx”])
File “/opt/homebrew/Caskroom/miniforge/base/envs/quantum/lib/python3.8/site-packages/pennylane/transforms/batch_transform.py”, line 371, in wrapper
_wrapper = self.qnode_wrapper(qnode, targs, tkwargs)
File “/opt/homebrew/Caskroom/miniforge/base/envs/quantum/lib/python3.8/site-packages/pennylane/transforms/metric_tensor.py”, line 319, in qnode_execution_wrapper
tkwargs.setdefault(“device_wires”, qnode.device.wires)
AttributeError: ‘torch.device’ object has no attribute ‘wires’

Does anyone know how can i define the custom optimizer for TorchLayer ?

Thank you :slight_smile:

Hi @Andre_Sequeira,

Thank you for asking this question here!
I’m not able to reproduce your exact error because I don’t have access to policy_estimator but I did manage to get a result by making the following changes:

Instead of having g = metric_tensor(p) I have g = metric_tensor(p).detach().numpy()

In the closure instead of having return loss, qml.metric_tensor(policy_estimator.qlayer) I have return loss, qml.metric_tensor(circuit)

I think the issue you were having was because qml.metric_tensor takes a tape, and policy_estimator.qlayer isn’t a tape.

Please let me know if this solves the issue you’re having! :smiley:

Hi @CatalinaAlbornoz,

thank you for your help. In my case policy_estimator is a class where i define a TorchLayer. Therefore, policy_estimator.qlayer is a TorchLayer as defined below.

class policy_estimator_q(nn.Module):        
    def __init__(self, env):
        super(policy_estimator_q, self).__init__()
        if policy == "Q":
            weight_shapes = {"weights":(n_layers, n_qubits, 2)}
            import functools

            if args.init == "random_0_2pi":
                #self.qlayer = qml.qnn.TorchLayer(qcircuit, weight_shapes)
                self.init_method = functools.partial(torch.nn.init.uniform_, a=0, b=2*np.pi)
                self.qlayer = qml.qnn.TorchLayer(qcircuit, weight_shapes, self.init_method)

Therefore i think i cannot pass to metric_tensor the circuit as you did because if i do that for instance accessing the qnode inside the torchlayer as qml.metric_tensor(policy_estimator.qlayer.qnode) then the metric_tensor is not able to reach the trainable parameters, i don’t know why. It says that weights are missing.
Any idea how to fix this ?
Thank you once again.

Hi @Andre_Sequeira,

Please excuse the delay in my response.
I have found some issues with your code, which I will break into issues with the optimizer, and other issues.

Issues with the optimizer:

  1. The g.shape[0] is causing some errors. However group["lam"] in the optimizer is always zero so you can actually bypass this error by commenting the line g += group["lam"] * np.identity(g.shape[0]).

  2. g and grad are Torch tensors so you can’t use them with Numpy directly. You need to use g.detach().numpy(), and grad.detach().numpy().

  3. g is a torch tensor which contains a single value so np.linalg.solve returns an error.

  4. group["lr"] is also constant, I’m not sure whether or not this is what you expected.

In any case these issues will require you re-thinking the optimizer.

Issues with the circuit

  1. Your circuit must include an argument called “inputs”, which corresponds to the non-trainable parameters. It can be None, but it needs to be an argument of your qnode. For more details you can check out the documentation for qnn.TorchLayer.
  2. Because of the previous item, you will need to change the way you call your circuit in the closure and anywhere else.

Issues in the closure

  1. Your policy_estimator was a bit complicated so I used return loss, qml.qnn.TorchLayer(circuit, weight_shapes) for the return of the closure, where weight_shapes = {"params":4} if you’re using the example circuit here.

Given the issues with the optimizer I managed to make it run but it doesn’t optimize well.

I hope this helps you move forward. Please let me know if any of this isn’t clear and I can share the full code that runs.

Also please let me know if this is what you were looking for!