hqm.layers.quanvolution

  1import torchvision
  2import torch
  3import sys
  4
  5sys.path += ['.', './utils/', '/circuits/']
  6
  7from hqm.circuits.circuit import QuantumCircuit
  8from hqm.utils.aiinterface import AIInterface
  9
 10
 11class Quanvolution2D(torch.nn.Module):
 12    '''
 13        Quanvolution2D layer.
 14
 15        Currently supports only Torch.
 16    '''
 17
 18    def __init__(self, qcircuit : QuantumCircuit, filters : int, kernelsize : int = 3, stride : int = 1, padding : str = 'same', aiframework : str = 'torch') -> None:
 19        '''
 20        Quanvolution2D constructor.  
 21
 22        Parameters:  
 23        -----------  
 24        - qcircuit : hqm.circuits.circuit.QuantumCircuit  
 25            QuantumCircuit object to be embedded into the quantum layer
 26        - filters : int
 27            number of quanvolution filters
 28        - kernelsize : int
 29            size of quanvolution kernel
 30        - stride : int
 31            stride for quanvolution operation
 32        - padding : str
 33            padding mode, same of valid
 34        - aiframework : str    
 35            string representing the AI framework in use, can be 'torch' or 'keras'. This will create  
 36            a compatible trainable layer for the framework.
 37
 38        Returns:    
 39        --------     
 40        Nothing, a Quanvolution2D object will be created.  
 41        '''
 42
 43        super().__init__()
 44
 45        if aiframework not in ['torch', 'keras']: raise Exception(f"Quanvolution2D curently supports only 'torch' as framework, found {aiframework}")
 46        if kernelsize < 1:                        raise Exception(f"kernelsize must be greater than 1, found {kernelsize}")
 47        if stride < 1:                            raise Exception(f"stride must be greater than 1, found {stride}")
 48        
 49        self.aiframework = aiframework
 50        self.n_qubits    = qcircuit.n_qubits
 51        
 52        if kernelsize**2 > self.n_qubits:         raise Exception(f"kernelsize**2 must be lower than n_qubits, found kernelsize**2={kernelsize**2} and {self.n_qubits}")
 53        if filters > self.n_qubits:               raise Exception(f"filters must be lower than n_qubits, found {filters} and {self.n_qubits}")
 54        
 55        self.filters    = filters 
 56        self.kernelsize = kernelsize
 57        self.stride     = stride
 58        self.padding    = padding
 59        self.qlayer     = AIInterface.network_layer(
 60                                circuit      = qcircuit.circuit, 
 61                                weight_shape = qcircuit.weight_shape, 
 62                                n_qubits     = qcircuit.n_qubits, 
 63                                aiframework  = self.aiframework
 64                            )
 65        
 66    def forward(self, x : torch.Tensor) -> torch.Tensor:
 67        '''
 68        Torch forward function for quanvolution layer
 69
 70        Parameters:
 71        -----------
 72        - x : torch.Tensor
 73            input image or tensor
 74        
 75        Returns:
 76        --------
 77        - out : torch.Tensor
 78            quanvoluted input
 79        '''
 80
 81        if len(x.shape) != 4: raise Exception(f"x must be a tensor of 4 elements (batch, channels, width, height), found {len(x.shape)}")
 82
 83        # Calculates the image shape after the convolution
 84        bs, ch, h, w = x.shape
 85        
 86        h_out = int(((h-self.kernelsize) / self.stride) +1)
 87        w_out = int(((w-self.kernelsize) / self.stride) +1)
 88        
 89        if self.padding == 'same':
 90            w_pad = int(round((w - w_out)/2))
 91            h_pad = int(round((h - h_out)/2))
 92        
 93        if self.padding == 'valid':
 94            h_pad = 0
 95            w_pad = 0 
 96
 97        out = torch.zeros((bs, self.filters, h_out, w_out, ch))
 98
 99        # Batch Loop
100        for b in range(bs):
101            # Channel Loop
102            for c in range(ch):
103                # Spatial Loops                                                
104                for j in range(0, h_out, self.stride):
105                    for k in range(0, w_out, self.stride):            
106                        # Process a kernel_size*kernel_size region of the images
107                        # with the quantum circuit stride*stride
108                        p = x[b, c, j:j+self.kernelsize, k:k+self.kernelsize].reshape(-1)
109                        q_results = self.qlayer(p)
110
111                        for f in range(self.filters):
112                            #out[b, f, j // self.kernelsize, k // self.kernelsize, c] = q_results[f]
113                            out[b, f, j:j+self.kernelsize, k:k+self.kernelsize, c] = q_results[f]
114
115        out = torch.mean(out, axis=-1)
116        
117        
118        if self.padding == 'same':
119            out = torch.nn.functional.pad(out, (h_pad, h_pad, w_pad, w_pad), "constant", 0)
120            out = torchvision.transforms.Resize([w,h])(out)
121
122        return out