Creating custom pytorch layers

Zohim_Chandani: <p>I am trying to create custom layers within PyTorch which can perform forward passes of data and run the backward method via finite difference. The custom layer is defined via a quantum circuit created in CUDA Quantum.</p>
<p>A simplified version of the code is here:</p>
<pre><code class="lang-auto"># !pip install cuda-quantum 


import cudaq 
from cudaq import spin 

import numpy as np 

import torch 
import torch.nn as nn

import torch
from torch.autograd import Function
import torch.nn as nn

cudaq.set_target("qpp-cpu")  
device = 'cpu'

torch.manual_seed(33)
np.random.seed(42)

</code></pre>
<pre><code class="lang-auto">class QuantumFunction(Function):
    
    """This class initializes the quantum circuit structure and defines the forward and the backward method."""

    def __init__(self, qubit_count, hamiltonian):
        
        """Define the quantum circuit in CUDA Quantum"""

        kernel, x, weights  = cudaq.make_kernel(list, list)  #x is the data input and weights are the parameters we want to optimise 

        self.kernel = kernel
        self.hamiltonian = hamiltonian

        qubits = kernel.qalloc(qubit_count)

        for i in range(qubit_count):
            kernel.ry(x[i] * weights[i], qubits[i])  #contains data*weights - similar to wx+b but without the bias 

        kernel.ry(x[0] * weights[0], qubits[0])  
        kernel.ry(x[1] * weights[1], qubits[1]) 
        
        kernel.rx(weights[2], qubits[0])
        kernel.rx(weights[3], qubits[1])

            
  
    def run(self, x_vals, weight_vals):
        
        """Execute the quantum circuit to output an expectation value"""
            
        N = x_vals.shape[0]  # Number of input samples  
    
        weight_vals = weight_vals.unsqueeze(0).repeat(N, 1)    #tile the weights so that we can run it for N data inputs 
        
        results = cudaq.observe(self.kernel, self.hamiltonian, x_vals, weight_vals).expectation()
        
        return results

    
    @staticmethod
    def forward(ctx, x_vals, weight_vals, quantum_circuit, shift):
        
        # Save shift and quantum_circuit in context to use in backward.
        ctx.shift = shift
        ctx.quantum_circuit = quantum_circuit

        # Calculate expectation value.
        exp_vals = ctx.quantum_circuit.run(x_vals, weight_vals)

        ctx.save_for_backward(x_vals, weight_vals)

        return exp_vals

    @staticmethod
    def backward(ctx, grad_output):
        """Backward pass computation via finite difference parameter shift"""

        x_vals, weight_vals = ctx.saved_tensors

        gradients = torch.zeros(len(weight_vals), device=device)

        for i in range(len(weight_vals)):
            shift_right = torch.clone(weight_vals)

            shift_right[i] += ctx.shift

            shift_left = torch.clone(weight_vals)

            shift_left[i] -= ctx.shift

            expectation_right = ctx.quantum_circuit.run(x_vals, shift_right)
            expectation_left = ctx.quantum_circuit.run(x_vals, shift_left)

            gradients[i] =  (expectation_right - expectation_left) / 2*ctx.shift
            
        
        return gradients * grad_output.float() , None, None
    
    
class QuantumLayer(nn.Module):
    """Encapsulates a quantum function into a quantum torch layer"""

    def __init__(self, qubit_count: int, hamiltonian, shift: torch.tensor):
        super(QuantumLayer, self).__init__()
        
        self.quantum_circuit = QuantumFunction(qubit_count , hamiltonian)  
        self.shift = shift

    def forward(self, input):
        ans = QuantumFunction.apply(input, self.quantum_circuit, self.shift)

        return ans
    
    
</code></pre>
<pre><code class="lang-auto">qubit_count = 2
hamiltonian = spin.z(0)
N = 128  # number of data samples 
latent_dim = 2
no_of_weights = 4
shift = torch.tensor(np.pi / 2)

x_vals = torch.rand(N, latent_dim)
weight_vals = torch.rand(no_of_weights)  #parameters we want to optimize 

generator = QuantumLayer(qubit_count, hamiltonian, shift)

generator_optimizer = torch.optim.Adam(generator.parameters(), lr=0.001)


</code></pre>
<pre><code class="lang-auto">
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[3], line 13
      9 weight_vals = torch.rand(no_of_weights)  #parameters we want to optimize 
     11 generator = QuantumLayer(qubit_count, hamiltonian, shift)
---&gt; 13 generator_optimizer = torch.optim.Adam(generator.parameters(), lr=0.001)

File ~/.local/lib/python3.10/site-packages/torch/optim/adam.py:45, in Adam.__init__(self, params, lr, betas, eps, weight_decay, amsgrad, foreach, maximize, capturable, differentiable, fused)
     39     raise ValueError(f"Invalid weight_decay value: {weight_decay}")
     41 defaults = dict(lr=lr, betas=betas, eps=eps,
     42                 weight_decay=weight_decay, amsgrad=amsgrad,
     43                 maximize=maximize, foreach=foreach, capturable=capturable,
     44                 differentiable=differentiable, fused=fused)
---&gt; 45 super().__init__(params, defaults)
     47 if fused:
     48     if differentiable:

File ~/.local/lib/python3.10/site-packages/torch/optim/optimizer.py:261, in Optimizer.__init__(self, params, defaults)
    259 param_groups = list(params)
    260 if len(param_groups) == 0:
--&gt; 261     raise ValueError("optimizer got an empty parameter list")
    262 if not isinstance(param_groups[0], dict):
    263     param_groups = [{'params': param_groups}]

ValueError: optimizer got an empty parameter list
</code></pre>
<p>Questions:</p>
<ol>
<li>I am unsure how to register the weights such that they can be picked up by the optimizer.</li>
<li>Does PyTorch have built-in finite difference backward methods?</li>
<li>Does PyTorch have built-in gradient-free methods such as COBYLA?</li>
</ol>
<p>Thanks.</p>

smth: <p>Please look at this tutorial: <a href="https://pytorch.org/tutorials/beginner/examples_nn/polynomial_module.html?highlight=nn%20parameter" class="inline-onebox">PyTorch: Custom nn Modules — PyTorch Tutorials 2.2.0+cu121 documentation</a></p>
<p>Basically, if you want a learnable parameter that can be seen by the Optimizer, you have to register it as a <code>nn.Parameter</code>. For example in your case, I think you want <code>self.shift = nn.Parameter(shift)</code> maybe?</p>
<blockquote>
<p>Does PyTorch have built-in finite difference backward methods?</p>
</blockquote>
<p>Only for unit-testing, its available in the utility <code>from torch.testing._internal.common_utils import gradcheck</code>. Generally, because we do partial derivatives, finite difference based backward becomes <code>O(number_of_parameters)</code>, which is not common in the kind of stuff PyTorch is optimized for (neural nets), but as you are writing custom autograd Functions, you can make the backward to just be finite differences yourself.</p>
<blockquote>
<ol start="2">
<li>Does PyTorch have built-in gradient-free methods such as COBYLA?</li>
</ol>
</blockquote>
<p>No. But there are packages such as <a href="https://github.com/facebookresearch/nevergrad">nevergrad</a> that implement gradient-free methods and are compatible with pytorch</p>

J_Johnson: <p>One other library built on PyTorch you might look into is PennyLane.</p>
<aside class="onebox allowlistedgeneric" data-onebox-src="https://docs.pennylane.ai/en/stable/introduction/pennylane.html">
  <header class="source">
      <img src="https://docs.pennylane.ai/en/stable/_static/favicon.ico" class="site-icon" width="" height="">

      <a href="https://docs.pennylane.ai/en/stable/introduction/pennylane.html" target="_blank" rel="noopener nofollow ugc">docs.pennylane.ai</a>
  </header>

  <article class="onebox-body">
    <img width="" height="" src="https://docs.pennylane.ai/en/stable/_static/header-tall.png" class="thumbnail">

<h3><a href="https://docs.pennylane.ai/en/stable/introduction/pennylane.html" target="_blank" rel="noopener nofollow ugc">What is PennyLane?</a></h3>

  <p>, PennyLane is a cross-platform Python library for programming quantum computers. Its differentiable programming paradigm enables the execution and training of quantum programs on various backends....</p>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>It is specifically designed for quantum programming and machine learning.</p>