Ai model on pytorch

Hello everyone. I’m trying to make my own model for deciphering Russian handwriting using cyrillic handwritten dataset for a school project.

import torch
import torchvision.transforms as transforms
from torchvision import datasets
import torch.nn as nn
from PIL import Image
import torch.nn.functional as F
import torch.optim as optim
import matplotlib.pyplot as plt
import math
from torch.utils.data import Dataset, DataLoader
import os
import pandas as pd

def preprocess_image(image_path):
    # Загрузка изображения и преобразование в оттенки серого
    image = Image.open(image_path).convert('L')

    # Определение последовательности преобразований
    transform = transforms.Compose([
        transforms.Resize((28, 28)),  # Изменение размера до 28x28 пикселей
        transforms.ToTensor(),  # Преобразование в тензор PyTorch
        transforms.Normalize((0.5,), (0.5,))  # Нормализация значений пикселей
    ])

    # Применение преобразований и добавление размерности пакета
    return transform(image).unsqueeze(0)


class HybridModel(nn.Module):
    def __init__(self):
        super(HybridModel, self).__init__()
        # Определите слои вашей модели здесь
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, stride=1, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
        self.fc1 = nn.Linear(16 * 64 * 64, 128)  # Примерный размер после свертки
        self.fc2 = nn.Linear(128, 10)  # Например, 10 классов

    def forward(self, x):
        # Определите, как данные проходят через слои
        x = self.pool(F.relu(self.conv1(x)))
        x = x.view(-1, 16 * 64 * 64)  # Измените форму для полносвязного слоя
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x  # Возвращаем выходные данные
class PrimaryCaps(nn.Module):
    def __init__(self, num_capsules=8, in_channels=256, out_channels=32):
        # Вызов конструктора родительского класса
        super(PrimaryCaps, self).__init__()

        # Создание списка из 8 сверточных слоев (капсул)
        # Каждая капсула преобразует 256 входных каналов в 32-мерный вектор выхода
        self.capsules = nn.ModuleList([
            nn.Conv2d(in_channels, out_channels, kernel_size=9, stride=2, padding=0)
            for _ in range(num_capsules)
        ])

    def forward(self, x):
        # Применение каждой капсулы к входным данным
        u = [capsule(x) for capsule in self.capsules]
        # Объединение результатов всех капсул
        u = torch.stack(u, dim=1)
        # Преобразование размерности для получения векторного представления
        u = u.view(x.size(0), 32 * 6 * 6, -1)
        # Применение функции squash для нормализации векторов
        return self.squash(u)

    def squash(self, input_tensor):
        # Вычисление квадрата нормы вектора
        squared_norm = (input_tensor ** 2).sum(-1, keepdim=True)
        # Вычисление коэффициента масштабирования
        scale = squared_norm / (1 + squared_norm)
        # Нормализация вектора с сохранением направления
        return scale * input_tensor / torch.sqrt(squared_norm)


class DigitCaps(nn.Module):
    def __init__(self, num_capsules=10, num_routes=32 * 6 * 6, in_channels=8, out_channels=16):
        # Вызов конструктора родительского класса
        super(DigitCaps, self).__init__()

        # Установка параметров слоя
        self.num_routes = num_routes  # Количество маршрутов от предыдущего слоя
        self.num_capsules = num_capsules  # Количество выходных капсул (по одной на каждую цифру)

        # Создание матрицы преобразования для динамической маршрутизации
        self.W = nn.Parameter(torch.randn(1, num_routes, num_capsules, out_channels, in_channels))

    def forward(self, x):
        # Получение размера пакета
        batch_size = x.size(0)

        # Добавление необходимых размерностей для матричных операций
        x = x.unsqueeze(2).unsqueeze(4)

        # Вычисление предсказаний для маршрутизации
        u_hat = torch.matmul(self.W, x).squeeze(4)

        # Инициализация логитов маршрутизации нулями
        b_ij = x.new_zeros(1, self.num_routes, self.num_capsules, 1)

        # Итеративный процесс динамической маршрутизации
        num_iterations = 3
        for iteration in range(num_iterations):
            # Вычисление коэффициентов связи через softmax
            c_ij = F.softmax(b_ij, dim=2)

            # Взвешенное суммирование предсказаний
            s_j = (c_ij * u_hat).sum(dim=1, keepdim=True)

            # Применение функции squash для получения выходных векторов
            v_j = self.squash(s_j)

            # Обновление логитов маршрутизации (кроме последней итерации)
            if iteration < num_iterations - 1:
                # Вычисление скалярного произведения для обновления весов
                a_ij = torch.matmul(u_hat.transpose(2, 3), v_j).squeeze(3).squeeze(2)
                # Обновление логитов маршрутизации
                b_ij = b_ij + a_ij.unsqueeze(2).unsqueeze(3)

        # Возврат финальных выходных векторов
        return v_j.squeeze(1)

    def squash(self, input_tensor):
        # Вычисление квадрата нормы вектора
        squared_norm = (input_tensor ** 2).sum(-1, keepdim=True)
        # Вычисление коэффициента масштабирования
        scale = squared_norm / (1 + squared_norm)
        # Нормализация вектора с сохранением направления
        return scale * input_tensor / torch.sqrt(squared_norm)


class Decoder(nn.Module):
    def __init__(self, num_classes=10, output_size=28):
        # Вызов конструктора родительского класса
        super(Decoder, self).__init__()

        # Создание последовательности слоев для реконструкции изображения
        self.reconstruct = nn.Sequential(
            # Первый полносвязный слой: преобразует капсульное представление в 512 нейронов
            nn.Linear(16 * num_classes, 512),
            nn.ReLU(inplace=True),  # Функция активации ReLU

            # Второй слой: расширяет представление до 1024 нейронов
            nn.Linear(512, 1024),
            nn.ReLU(inplace=True),

            # Выходной слой: формирует изображение размером output_size × output_size
            nn.Linear(1024, output_size * output_size),
            nn.Sigmoid()  # Сигмоида для нормализации значений пикселей в диапазоне [0,1]
        )

    def forward(self, x, data):
        # Вычисление длины векторов капсул (вероятностей классов)
        classes = (x ** 2).sum(dim=-1) ** 0.5
        # Применение softmax для получения распределения вероятностей
        classes = F.softmax(classes, dim=-1)

        # Определение класса с максимальной вероятностью
        _, max_length_indices = classes.max(dim=1)
        # Маскирование входа: оставляем только активации для предсказанного класса
        masked = x * F.one_hot(max_length_indices, num_classes=10).unsqueeze(-1)

        # Реконструкция изображения
        reconstructions = self.reconstruct(masked.view(x.size(0), -1))
        # Преобразование выхода в формат изображения (batch_size, 1, 28, 28)
        reconstructions = reconstructions.view(-1, 1, 28, 28)

        # Возвращаем реконструированное изображение и вероятности классов
        return reconstructions, classes


def loss_function(reconstructions, classes, images, labels, lambda_recon=0.0005):
    # Вычисляем кросс-энтропийную потерю для задачи классификации
    margin_loss = F.cross_entropy(classes, labels)

    # Вычисляем среднеквадратичную ошибку для задачи реконструкции
    reconstruction_loss = F.mse_loss(reconstructions, images)

    # Возвращаем комбинированную функцию потерь
    return margin_loss + lambda_recon * reconstruction_loss


def train(model, train_loader, optimizer, epoch):
    model.train()  # Переводим модель в режим обучения
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()  # Обнуляем градиенты

        # Прямой проход через модель
        reconstructions, classes = model(data)

        # Вычисляем функцию потерь
        # Предполагается, что ваша функция потерь принимает реконструкции, классы, данные и цели
        loss = loss_function(reconstructions, classes, data, target)

        loss.backward()  # Обратное распространение
        optimizer.step()  # Обновляем параметры модели

        # Выводим информацию о ходе обучения каждые 100 батчей
        if batch_idx % 100 == 0:
            print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} '
                  f'({100. * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}')


def augment_capsules(capsules, num_augmentations=5):
    # Получаем размерность входных капсул: размер батча, количество классов и размер капсулы
    batch_size, num_classes, capsule_dim = capsules.size()
    augmented_capsules = []  # Список для хранения аугментированных капсул

    for _ in range(num_augmentations):
        # Случайное изменение длины вектора капсулы
        # Генерируем случайные коэффициенты масштабирования от 0.9 до 1.1
        scale = torch.rand(batch_size, num_classes, 1) * 0.2 + 0.9
        scaled_capsules = capsules * scale  # Масштабируем капсулы

        # Случайное вращение вектора капсулы
        # Генерируем случайные углы вращения
        theta = torch.rand(batch_size, num_classes, 1) * 2 * math.pi
        # Создаём матрицу вращения из углов
        rotation_matrix = torch.cat([torch.cos(theta), -torch.sin(theta),
                                     torch.sin(theta), torch.cos(theta)], dim=-1)
        rotation_matrix = rotation_matrix.view(batch_size, num_classes, 2, 2)  # Преобразуем в 2x2 матрицу

        # Применяем вращение к масштабированным капсулам
        rotated_capsules = torch.matmul(rotation_matrix, scaled_capsules.view(batch_size, num_classes, 2, -1))
        rotated_capsules = rotated_capsules.view(batch_size, num_classes, -1)  # Возвращаем к исходной размерности

        # Добавляем аугментированные капсулы в список
        augmented_capsules.append(rotated_capsules)

    # Объединяем оригинальные капсулы с аугментированными
    return torch.cat([capsules] + augmented_capsules, dim=0)


def test(model, test_loader):
    model.eval()  # Переводим модель в режим оценки
    test_loss = 0  # Инициализируем переменную для накопления потерь
    correct = 0  # Инициализируем счётчик правильных предсказаний
    with torch.no_grad():  # Отключаем вычисление градиентов для экономии памяти
        for data, target in test_loader:
            reconstructions, classes = model(data)  # Выполняем прямой проход через модель
            # Вычисляем функцию потерь
            test_loss += loss_function(reconstructions, classes, data, target).item()
            # Получаем предсказания модели
            pred = classes.argmax(dim=1, keepdim=True)
            # Сравниваем предсказания с истинными метками
            correct += pred.eq(target.view_as(pred)).sum().item()

    # Вычисляем среднюю потерю и точность
    test_loss /= len(test_loader.dataset)
    accuracy = 100. * correct / len(test_loader.dataset)

    # Выводим результаты оценки
    print(f'\nTest set: Average loss: {test_loss:.4f}, '
          f'Accuracy: {correct}/{len(test_loader.dataset)} ({accuracy:.2f}%)\n')

    return test_loss, accuracy


def visualize_reconstructions(model, data_loader, num_images=5):
    model.eval()  # Переводим модель в режим оценки
    with torch.no_grad():  # Отключаем вычисление градиентов для экономии памяти
        for data, _ in data_loader:
            reconstructions, _ = model(data)  # Получаем реконструкции от модели
            break  # Выходим из цикла после получения первого батча

    # Создаём фигуру для визуализации
    fig, axes = plt.subplots(2, num_images, figsize=(12, 6))
    for i in range(num_images):
        # Отображаем оригинальные изображения
        axes[0, i].imshow(data[i].squeeze(), cmap='gray')
        axes[0, i].axis('off')  # Убираем оси
        # Отображаем реконструированные изображения
        axes[1, i].imshow(reconstructions[i].squeeze(), cmap='gray')
        axes[1, i].axis('off')  # Убираем оси

    plt.tight_layout()  # Убираем лишние отступы
    plt.show()  # Показываем визуализацию


class CustomDataset(Dataset):
    def __init__(self, folder_path, labels_file, transform=None):
        self.folder_path = folder_path
        self.transform = transform

        # Загружаем метки из файла TSV, указывая разделитель
        self.labels = pd.read_csv(labels_file, sep='\t', header=None, names=['filename', 'label'])

        # Извлекаем имена файлов и метки из DataFrame
        self.images = self.labels['filename'].tolist()  # Список имен файлов изображений
        self.targets = self.labels['label'].tolist()  # Список меток

    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        img_name = os.path.join(self.folder_path, self.images[idx])
        image = Image.open(img_name).convert('RGB')

        if self.transform:
            image = self.transform(image)

        label = self.targets[idx]  # Получаем метку для текущего изображения
        return image, label  # Возвращаем изображение и метку

# Примените необходимые преобразования
transform = transforms.Compose([
    transforms.Resize((128, 128)),
    transforms.ToTensor(),
])

# Загрузка обучающих данных
train_dataset = CustomDataset(
    folder_path='C:/Users/Semka/PycharmProjects/RuChif/train',
    labels_file='C:/Users/Semka/PycharmProjects/RuChif/train.tsv',  # Путь к файлу меток
    transform=transform
)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

# Загрузка тестовых данных
test_dataset = CustomDataset(
    folder_path='C:/Users/Semka/PycharmProjects/RuChif/test',
    labels_file='C:/Users/Semka/PycharmProjects/RuChif/test.tsv',  # Путь к файлу меток для теста
    transform=transform
)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)

# Инициализация модели и оптимизатора
model = HybridModel() #модель
optimizer = optim.Adam(model.parameters())  # Инициализируем оптимизатор Adam

# Обучение модели
num_epochs = 10  # Определите количество эпох
for epoch in range(1, num_epochs + 1):  # Цикл по эпохам
    train(model, train_loader, optimizer, epoch)  # Запускаем процесс обучения
    # Сохраняем модель после каждой эпохи
    torch.save(model.state_dict(), f'model_epoch_{epoch}.pth')

# Оценка модели после обучения
test_loss, test_accuracy = test(model, test_loader)

# Визуализация реконструкций
visualize_reconstructions(model, test_loader)

but an error pops up in my code.

Traceback (most recent call last):
  File "C:\Users\Semka\PycharmProjects\RuChif\all test.py", line 334, in <module>
    train(model, train_loader, optimizer, epoch)  # Запускаем процесс обучения
  File "C:\Users\Semka\PycharmProjects\RuChif\all test.py", line 189, in train
    reconstructions, classes = model(data)
ValueError: too many values to unpack (expected 2)

please help someone, I will be glad of any help.