hqm.classification.hcnn

 1import sys
 2sys.path += ['.', './layers/', './utils/']
 3
 4from hqm.layers.basiclayer import BasicLayer
 5from hqm.utils.sizes import size_conv_layer
 6
 7import torch
 8
 9class HybridLeNet5(torch.nn.Module):
10    '''
11        This class implements a quantum hybrid convolutional neural network based on LeNet-5.
12        HybridLeNet5 is composed of classical convlutional block and hybrid quantum MLP.
13        The size of the network output is defined by ou_dim.
14    '''
15
16    def __init__(self, qlayer : BasicLayer, in_shape : tuple, ou_dim : int) -> None:
17        '''
18            HybridLeNet5 constructor.  
19
20            Parameters:  
21            -----------  
22            - qlayer : hqm.layers.basilayer.BasicLayer  
23                hqm quantum layer to be stacked between two fully connected layers  
24            - in_shape : tuple  
25                tuple representing the shape of the input image (channels, widht, height)  
26            - ou_dim : int  
27                integer representing the output size of the hybrid model  
28            
29            Returns:  
30            --------  
31            Nothing, a HybridLeNet5 object will be created.    
32        '''
33
34        super().__init__()
35
36        if len(in_shape) != 3: raise Exception(f"The parameter in_shape must be a tuple of four elements (channels, widht, height), found {in_shape}")
37        if ou_dim < 1: raise Exception(f"The parameter ou_dim must be greater than 1, found {ou_dim}")
38        
39        c, w, h = in_shape
40        
41        c1 = 6
42        self.conv_1    = torch.nn.Conv2d(in_channels=c, out_channels=c1, kernel_size=5, padding=2, stride=1)
43        w1 = size_conv_layer(w, kernel_size=5, padding=2, stride=1)
44        h1 = size_conv_layer(h, kernel_size=5, padding=2, stride=1)
45        
46        self.max_pool1 = torch.nn.MaxPool2d(kernel_size = (2,2), stride=(2,2))
47        w2 = size_conv_layer(w1, kernel_size=2, padding=0, stride=2)
48        h2 = size_conv_layer(h1, kernel_size=2, padding=0, stride=2)
49        
50        c2 = 16
51        self.conv_2  = torch.nn.Conv2d(in_channels=c1, out_channels=c2, kernel_size=5,  stride=1)
52        w3 = size_conv_layer(w2, kernel_size=5, padding=0, stride=1)
53        h3 = size_conv_layer(h2, kernel_size=5, padding=0, stride=1)
54
55        self.max_pool2 = torch.nn.MaxPool2d(kernel_size = (2,2), stride=(2,2))
56        w4 = size_conv_layer(w3, kernel_size=2, padding=0, stride=2)
57        h4 = size_conv_layer(h3, kernel_size=2, padding=0, stride=2)
58
59        self.flatten_size = w4 * h4 * c2
60        fc_2_size = int(self.flatten_size * 30 / 100)
61
62        self.fc_1    = torch.nn.Linear(self.flatten_size, fc_2_size)
63        self.fc_2    = torch.nn.Linear(fc_2_size, qlayer.n_qubits)
64        self.qc_1    = qlayer.qlayer
65        self.fc_3    = torch.nn.Linear(qlayer.n_qubits, ou_dim)
66        self.relu    = torch.nn.ReLU()
67        self.softmax = torch.nn.Softmax(dim=1)
68    
69    def forward(self, x : torch.Tensor) -> torch.Tensor:
70        '''
71            Torch forward method  
72
73            Parameters:  
74            -----------
75            - x : torch.Tensor  
76                input for the torch model  
77
78            Returns:  
79            --------  
80            - out : torch.Tensor  
81                output from the torch model  
82        '''
83        
84        x = self.max_pool1(self.relu(self.conv_1(x)))
85        x = self.max_pool2(self.relu(self.conv_2(x)))
86        x = x.view(-1, self.flatten_size)
87        x = self.relu(self.fc_1(x))
88        x = self.relu(self.fc_2(x))
89        x = self.relu(self.qc_1(x))
90        x = self.fc_3(x)
91        out = self.softmax(x)
92        return out