classification/classifier.py - master - IvlevAlexander0/QNN_FastApi

classification/classifier.py
226 строк9 KB
IvlevAlexander0
ver0.9.2
24 окт 2025, 10:47
24 окт 2025, 10:471eea02a
import numpy as np

import torch
import torch.nn as nn
from torchvision.transforms import v2

from qiskit import QuantumCircuit
from qiskit.circuit import ParameterVector
from qiskit.circuit.library import RealAmplitudes
from qiskit_machine_learning.neural_networks import SamplerQNN
from qiskit_machine_learning.connectors import TorchConnector
from qiskit_machine_learning.utils import algorithm_globals


class QuantumHybridModel(nn.Module):
    """ Гибридная квантово-классическая нейронная сеть """
    def __init__(self, config):
        '''
        config: набор параметров для создания модели
        '''
        super(QuantumHybridModel, self).__init__()
        self.num_classes = config['num_classes']
        self.num_qubits = config['num_qubits']
        self.quantum_dim = config['quantum_dim']
        self.output_dim = config['output_dim']
        if(self.output_dim == None):
            self.output_dim = 2**self.num_classes

        # Классическая часть - ResNet
        self.classical_backbone = torch.hub.load('pytorch/vision', 'resnet18', pretrained=True)
        self.classical_backbone.fc = nn.Linear(512, self.quantum_dim*self.num_qubits)

        # Адаптер между классической и квантовой частью
        self.classical_to_quantum = nn.Sequential(
            nn.Softmax()
        )

        # Адаптер выхода квантовой части
        self.quantum_to_output = nn.Sequential(
            nn.Linear(2*self.output_dim*self.quantum_dim, self.num_classes)
        )

        # Квантовый слой
        self.quantum_layer = self._create_quantum_layer(self.num_qubits, self.output_dim)

        # Параллельный полносвязный слой
        self.sub_linear = nn.Sequential(
            nn.Linear(self.num_qubits, self.output_dim*self.quantum_dim),
            nn.ReLU()
        )

    def _create_init_circuit(self, num_qubits):
        """Создаёт схему загрузки классических данных в кубиты угловым кодированием"""
        '''
        Param:
            num_qubits: количество кубитов

        Return:
            qc: квантовая схема загрузки классических данных в кубиты
        '''
        qc = QuantumCircuit(num_qubits)
        pv = ParameterVector(name="init", length=num_qubits)
        for i in range(num_qubits):
            qc.ry(pv[i], i)
        return qc

    def _create_quantum_layer(self, num_qubits, output_dim = None):
        """Создаёт квантовый слой"""
        '''
        Param:
            num_qubits: количество кубитов;
            output_dim: размерность выхода квантового слоя

        Return:
            quantum_layer: квантовый слой в обёртке TorchConnector
        '''
        qc = QuantumCircuit(num_qubits)

        # feature_map - часть загрузки классических данных
        feature_map = self._create_init_circuit(num_qubits)
        qc.compose(feature_map, inplace=True)

        # VQC - обучаемая часть
        conv_size = int(np.log2(output_dim))
        ansatz = RealAmplitudes(num_qubits, reps=num_qubits-conv_size+1, skip_final_rotation_layer=True)
        qc.compose(ansatz.decompose(), inplace=True)
        ansatz2 = RealAmplitudes(conv_size, reps=conv_size, parameter_prefix="θ2", entanglement="full")
        qc.compose(ansatz2.decompose(), inplace=True)
        params_ans = ansatz.parameters.data
        for i in range(len(ansatz2.parameters.data)):
            params_ans.append(ansatz2.parameters.data[i])

        # Измеряем все кубиты
        qc.measure_all()

        # Создаём QNN
        def parity(x):
            if(output_dim != None):
                return x % output_dim
            else:
                return x

        interpret = None
        if(output_dim != None):
            interpret = parity

        qnn = SamplerQNN(
            sampler=None,
            circuit=qc,
            input_params=feature_map.parameters,
            weight_params=params_ans,
            input_gradients=True,
            pass_manager=None,
            output_shape=output_dim,
            interpret=interpret
        )

        # Обёртка для интеграции с PyTorch
        algorithm_globals.random_seed = 42
        quantum_layer = TorchConnector(qnn, 0.1 * (2 * algorithm_globals.random.random(qnn.num_weights) - 1))

        return quantum_layer

    def forward(self, image):
        """ Прямой проход модели"""
        '''
        Param:
            image: предобработанное изображение

        Return:
            output: тензор вероятностей пренадлежности классам
        '''
        # Классическая обработка
        classical_features = self.classical_backbone(image)

        # Подготовка для квантового слоя
        quantum_input = self.classical_to_quantum(classical_features)

        # Квантовая обработкая
        quantum_output = self.quantum_layer(quantum_input)

        # Параллельная классическая обработка
        sub_linear_output = self.sub_linear(quantum_input)

        # Финальная классификация
        output = self.quantum_to_output(torch.cat((quantum_output, sub_linear_output), dim=-1))

        return output

def return_model(config):
    """Создаёт и возвращает гибридную квантово-классическую модель"""
    '''
        Param:
            config: набор параметров для создания модели

        Return:
            model: гибридная квантово-классическая модель
    '''
    model = QuantumHybridModel(config)

    model.load_state_dict(torch.load(config['classifier_model'], map_location=config['device']))
    model = model.to(config['device'])
    model.eval()

    return model

def prediction(image, model, in_size, device, classes_indexes):
    """Функция классификации"""
    '''
        Param:
            image: изображение;
            model: модель нейронной сети;
            in_size: размер избражения подаваемого на вход модели (3ximage_sizeximage_size);
            device: устройство исполнения (cuda или cpu);
            classes_indices: словарь соответствия индексов классов и их названия

        Return:
            pred: предсказанный
    '''
    # Предобработка изображения
    data_transform = v2.Compose([
        v2.Resize((in_size, in_size)),
        v2.ToImage(), 
        v2.ToDtype(torch.float32, scale=True),
        v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
    image = data_transform(image).to(device)

    # Классификация
    output = model(image.unsqueeze(0))
    pred = classes_indexes[str(output.argmax(dim=1, keepdim=True).item())]
    return pred


def return_classifier(config):
    """Возвращает лямбда функцию классификации одного изображения"""
    '''
        Param:
            config: набор внешних параметров для создания модели:
                'device': Устройство исполнения (cuda или cpu)
                'image_size': Размер изoбражения подаваемого на вход модели (3ximage_sizeximage_size);
                'classifier_model': # Путь к сохранённой модели

        Return:
            Лямбда функция классификации одного изображения
    '''
    model_base = {
        'num_qubits': 5,    # Количество кубитов в слое
        'num_classes': 5,   # Количество классов 
        'quantum_dim': 1,   # Количество параллельных квантовых слоёв
        'output_dim': 8,    # Размерность выхода квантового слоя
    }

    config.update(model_base)

    model = return_model(config)

    classes_indexes = {
        "0": "Brick_defect",
        "1": "Ceramic_defect",
        "2": "Noncrack_free",
        "3": "Road_defect",
        "4": "Wall_defect"
    }

    return lambda img: prediction(img, model=model, in_size=config['image_size'], device=config['device'], classes_indexes=classes_indexes)