# This code is part of a Qiskit project.
#
# (C) Copyright IBM 2021, 2023.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE.txt file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.
"""A connector to use Qiskit (Quantum) Neural Networks as PyTorch modules."""
from __future__ import annotations
from typing import Tuple, Any, cast
import numpy as np
import qiskit_machine_learning.optionals as _optionals
from ..exceptions import QiskitMachineLearningError
from ..neural_networks import NeuralNetwork
if _optionals.HAS_TORCH:
import torch
# imports for inheritance and type hints
from torch import Tensor
from torch.autograd import Function
from torch.nn import Module
else:
class Function: # type: ignore
"""Empty Function class
Replacement if torch.autograd.Function is not present.
"""
pass
class Tensor: # type: ignore
"""Empty Tensor class
Replacement if torch.Tensor is not present.
"""
pass
class Module: # type: ignore
"""Empty Module class
Replacement if torch.nn.Module is not present.
"""
pass
[docs]@_optionals.HAS_TORCH.require_in_instance
class TorchConnector(Module):
"""Connects a Qiskit (Quantum) Neural Network to PyTorch."""
# pylint: disable=abstract-method
class _TorchNNFunction(Function):
# pylint: disable=arguments-differ
@staticmethod
def forward( # type: ignore
ctx: Any,
input_data: Tensor,
weights: Tensor,
neural_network: NeuralNetwork,
sparse: bool,
) -> Tensor:
"""Forward pass computation.
Args:
ctx: The context to be passed to the backward pass.
input_data: The input data.
weights: The weights.
neural_network: The neural network to be connected.
sparse: Indicates whether to use sparse output or not.
Returns:
The resulting value of the forward pass.
Raises:
QiskitMachineLearningError: Invalid input data.
RuntimeError: if connector is configured as sparse and the network is not sparse.
"""
# validate input shape
if input_data.shape[-1] != neural_network.num_inputs:
raise QiskitMachineLearningError(
f"Invalid input dimension! Received {input_data.shape} and "
+ f"expected input compatible to {neural_network.num_inputs}"
)
ctx.neural_network = neural_network
ctx.sparse = sparse
ctx.save_for_backward(input_data, weights)
# Detach the tensors and move it to CPU as we need numpy array to compute gradients
# of the quantum neural network. If the tensors are on CPU already this does nothing.
# Some other tensors down below are also moved to CPU for computations.
result = neural_network.forward(
input_data.detach().cpu().numpy(), weights.detach().cpu().numpy()
)
if ctx.sparse:
if neural_network.sparse:
_optionals.HAS_SPARSE.require_now("SparseArray")
# pylint: disable=import-error
from sparse import SparseArray, COO
# todo: replace output type from DOK to COO?
result = cast(COO, cast(SparseArray, result).asformat("coo"))
result_tensor = torch.sparse_coo_tensor(result.coords, result.data)
else:
raise RuntimeError(
"TorchConnector configured as sparse, the network must be sparse as well"
)
else:
# connector is dense
if neural_network.sparse:
# convert to dense
_optionals.HAS_SPARSE.require_now("SparseArray")
from sparse import SparseArray
# cast is required by mypy
result = cast(SparseArray, result).todense()
result_tensor = torch.as_tensor(result, dtype=torch.float)
# if the input was not a batch, then remove the batch-dimension from the result,
# since the neural network will always treat input as a batch and cast to a
# single-element batch if no batch is given and PyTorch does not follow this
# convention.
if len(input_data.shape) == 1:
result_tensor = result_tensor[0]
# place the resulting tensor back to the device where input data is stored
result_tensor = result_tensor.to(input_data.device)
return result_tensor
@staticmethod
def backward(ctx: Any, grad_output: Tensor) -> Tuple: # type: ignore
"""Backward pass computation.
Args:
ctx: context
grad_output: previous gradient
Raises:
QiskitMachineLearningError: Invalid input data.
RuntimeError: if connector is configured as sparse and the network is not sparse.
Returns:
gradients for the first two arguments and None for the others
"""
# get context data
input_data, weights = ctx.saved_tensors
neural_network = ctx.neural_network
# validate input shape
if input_data.shape[-1] != neural_network.num_inputs:
raise QiskitMachineLearningError(
f"Invalid input dimension! Received {input_data.shape} and "
+ f" expected input compatible to {neural_network.num_inputs}"
)
# ensure same shape for single observations and batch mode
if len(grad_output.shape) == 1:
grad_output = grad_output.view(1, -1)
# evaluate QNN gradient
input_grad, weights_grad = neural_network.backward(
input_data.detach().cpu().numpy(), weights.detach().cpu().numpy()
)
if input_grad is not None:
if ctx.sparse:
if neural_network.sparse:
_optionals.HAS_SPARSE.require_now("Sparse")
import sparse
from sparse import COO
grad_output = grad_output.detach().cpu()
grad_coo = COO(grad_output.indices(), grad_output.values())
# Takes gradients from previous layer in backward pass (i.e. later layer in
# forward pass) j for each observation i in the batch. Multiplies this with
# the gradient from this point on backwards with respect to each input k.
# Sums over all j to get total gradient of output w.r.t. each input k and
# batch index i. This operation should preserve the batch dimension to be
# able to do back-prop in a batched manner.
# Pytorch does not support sparse einsum, so we rely on Sparse.
# pylint: disable=no-member
input_grad = sparse.einsum("ij,ijk->ik", grad_coo, input_grad)
# return sparse gradients
input_grad = torch.sparse_coo_tensor(input_grad.coords, input_grad.data)
else:
# this exception should never happen
raise RuntimeError(
"TorchConnector configured as sparse, "
"the network must be sparse as well"
)
else:
# connector is dense
if neural_network.sparse:
# convert to dense
input_grad = input_grad.todense()
input_grad = torch.as_tensor(input_grad, dtype=torch.float)
# same as above
input_grad = torch.einsum("ij,ijk->ik", grad_output.detach().cpu(), input_grad)
# place the resulting tensor to the device where they were stored
input_grad = input_grad.to(input_data.device)
if weights_grad is not None:
if ctx.sparse:
if neural_network.sparse:
import sparse
from sparse import COO
grad_output = grad_output.detach().cpu()
grad_coo = COO(grad_output.indices(), grad_output.values())
# Takes gradients from previous layer in backward pass (i.e. later layer in
# forward pass) j for each observation i in the batch. Multiplies this with
# the gradient from this point on backwards with respect to each
# parameter k. Sums over all i and j to get total gradient of output
# w.r.t. each parameter k. The weights' dimension is independent of the
# batch size.
# pylint: disable=no-member
weights_grad = sparse.einsum("ij,ijk->k", grad_coo, weights_grad)
# return sparse gradients
weights_grad = torch.sparse_coo_tensor(
weights_grad.coords, weights_grad.data
)
else:
# this exception should never happen
raise RuntimeError(
"TorchConnector configured as sparse, "
"the network must be sparse as well"
)
else:
if neural_network.sparse:
# convert to dense
weights_grad = weights_grad.todense()
weights_grad = torch.as_tensor(weights_grad, dtype=torch.float)
# same as above
weights_grad = torch.einsum(
"ij,ijk->k", grad_output.detach().cpu(), weights_grad
)
# place the resulting tensor to the device where they were stored
weights_grad = weights_grad.to(weights.device)
# return gradients for the first two arguments and None for the others (i.e. qnn/sparse)
return input_grad, weights_grad, None, None
def __init__(
self,
neural_network: NeuralNetwork,
initial_weights: np.ndarray | Tensor | None = None,
sparse: bool | None = None,
):
"""
Args:
neural_network: The neural network to be connected to PyTorch. Remember
that ``input_gradients`` must be set to ``True`` in the neural network
initialization before passing it to the ``TorchConnector`` for the gradient
computations to work properly during training.
initial_weights: The initial weights to start training the network. If this is None,
the initial weights are chosen uniformly at random from [-1, 1].
sparse: Whether this connector should return sparse output or not. If sparse is set
to None, then the setting from the given neural network is used. Note that sparse
output is only returned if the underlying neural network also returns sparse output,
otherwise an error will be raised.
Raises:
QiskitMachineLearningError: If the connector is configured as sparse and the underlying
network is not sparse.
"""
super().__init__()
self._neural_network = neural_network
if sparse is None:
sparse = self._neural_network.sparse
self._sparse = sparse
if self._sparse and not self._neural_network.sparse:
# connector is sparse while the underlying neural network is not
raise QiskitMachineLearningError(
"TorchConnector configured as sparse, the network must be sparse as well"
)
weight_param = torch.nn.Parameter(torch.zeros(neural_network.num_weights))
# Register param. in graph following PyTorch naming convention
self.register_parameter("weight", weight_param)
# If `weight_param` is assigned to `self._weights` after registration,
# it will not be re-registered, and we can keep the private var. name
# "_weights" for compatibility. The alternative, doing:
# `self._weights = TorchParam(Tensor(neural_network.num_weights))`
# would register the parameter with the name "_weights".
self._weights = weight_param
if initial_weights is None:
self._weights.data.uniform_(-1, 1)
else:
self._weights.data = torch.tensor(initial_weights, dtype=torch.float)
@property
def neural_network(self) -> NeuralNetwork:
"""Returns the underlying neural network."""
return self._neural_network
@property
def weight(self) -> Tensor:
"""Returns the weights of the underlying network."""
return self._weights
@property
def sparse(self) -> bool | None:
"""Returns whether this connector returns sparse output or not."""
return self._sparse
[docs] def forward(self, input_data: Tensor | None = None) -> Tensor:
"""Forward pass.
Args:
input_data: data to be evaluated.
Returns:
Result of forward pass of this model.
"""
input_ = input_data if input_data is not None else torch.zeros(0)
return TorchConnector._TorchNNFunction.apply(
input_, self._weights, self._neural_network, self._sparse
)