Как распознать объекты на этом изображении?

У меня есть изображения, пример. Но у всех изображений размер изменён на 120x80. Мне нужно распознать, что находится на изображении: цифра (от 1 до 9) или буква (полный английский алфавит). Но моя модель не обучается. Она просто останавливается на ошибке ~3,6 (CrossEntropyLoss, 35 классов).

Затем я попытался увидеть, какое изображение выводятся после каждого слоя, и они после 3 блока (см. модель ниже) абсолютно одинаковы (есть единичные исключения), остаётся только белый фон. Мои объекты (цифры / буквы) не попадают в следующий слой. Я пробовал увеличить размер ядра Conv2d, уменьшить количество фильтров, но это не сработало.

Изменино: Использую pytorch. При обучении задействован оптимизатор Adam с lr = 0.001, batch_size пробовал что 32, что 64 - ни то, ни то не катит. Датасет разделен 20% - валидационная выборка, 80% - тренировочная. Пытался тренировать 100 и 500 эпох, результат один: график обучерия (синим - тренировочная выборка, жёлтым - валидационая).

Код для обучения:

from torch.utils.data import DataLoader, random_split
from torch.nn import CrossEntropyLoss
from torch.optim import Adam
from torchvision import transforms
import matplotlib.pyplot as plt

from dataset.dataset import CellsDataset
from model import RecognitionModel

batch_size = 32
epochs = 100
lr = 0.001

transform = transforms.Compose([
    transforms.Resize((80, 120)),  # Изменение размера изображений
    # transforms.RandomHorizontalFlip(),  # Случайное горизонтальное отражение
    # transforms.RandomRotation(20),  # Случайное вращение на 20 градусов
    # transforms.RandomAffine(degrees=15, translate=(0.1, 0.1)),  # Случайная аффинная трансформация
    transforms.Grayscale(),
    transforms.ToTensor(),  # Преобразование в тензор
])

dataset = CellsDataset(transform)

train_dataset, valid_dataset = random_split(dataset, [0.8, 0.2])
train_dataloader = DataLoader(train_dataset, batch_size = batch_size, shuffle = True)
valid_dataloader = DataLoader(valid_dataset, batch_size = batch_size, shuffle = True)

model = RecognitionModel()

loss_func = CrossEntropyLoss()
optimizer = Adam(model.parameters(), lr = lr)

train_losses = []
valid_losses = []

for i in range(epochs):
    # Train
    optimizer.zero_grad()
    image, label = next(iter(train_dataloader))

    pred = model(image)
    loss = loss_func(pred, label)
    train_losses.append(loss.item())

    # Validation
    image, label = next(iter(valid_dataloader))

    pred = model(image)
    loss_ = loss_func(pred, label)
    valid_losses.append(loss_.item())

    # Backward

    loss.backward()
    optimizer.step()

    print(f"Epoch {i+1}/{epochs} Loss {loss.item()} Validation loss {loss_.item()}")

Код модели:

import torch
from torch.nn import Module, Conv2d, MaxPool2d, ReLU, AdaptiveMaxPool2d, Linear, LeakyReLU, Softmax

class CNNBlock(Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv = Conv2d(
            in_channels = in_channels,
            out_channels = out_channels,
            kernel_size = 5,
            padding = "same"
        )
        self.act = ReLU()
        self.maxpool = MaxPool2d(
            kernel_size = 2,
            stride = 2
        )

    def forward(self, x):
        return self.maxpool(self.act(self.conv(x)))

class RecognitionModel(Module):
    def __init__(self):
        super().__init__()
        
        self.block1 = CNNBlock(1, 32)
        self.block2 = CNNBlock(32, 64)
        self.block3 = CNNBlock(64, 128)
        self.block4 = CNNBlock(128, 256)
        self.conv1 = Conv2d(
            in_channels = 256,
            out_channels = 512,
            kernel_size = 3
        )
        self.act1 = ReLU()
        self.conv2 = Conv2d(
            in_channels = 512,
            out_channels = 1024,
            kernel_size = 3
        )
        self.act2 = ReLU()
        self.globalmaxpool = AdaptiveMaxPool2d(output_size = 1)

        self.sqz = lambda x: x.squeeze()
        self.linear1 = Linear(
            in_features = 1024,
            out_features = 512
        )
        self.act3 = LeakyReLU()
        self.linear2 = Linear(
            in_features = 512,
            out_features = 256
        )
        self.act4 = LeakyReLU()
        self.linear3 = Linear(
            in_features = 256,
            out_features = 128
        )
        self.act5 = LeakyReLU()
        self.linear4 = Linear(
            in_features = 128,
            out_features = 64
        )
        self.act6 = LeakyReLU()
        self.linear5 = Linear(
            in_features = 64,
            out_features = 35
        )
        self.act7 = Softmax()

    def forward(self, x):
        x = self.block1(x)
        x = self.block2(x)
        x = self.block3(x)
        x = self.block4(x)
        x = self.act1(self.conv1(x))
        x = self.act2(self.conv2(x))
        x = self.globalmaxpool(x)
        x = self.sqz(x)
        x = self.act3(self.linear1(x))
        x = self.act4(self.linear2(x))
        x = self.act5(self.linear3(x))
        x = self.act6(self.linear4(x))
        y = self.act7(self.linear5(x))

        return y

Код для сборки датасета:

import json
import torch
from torch.utils.data import Dataset
from PIL import Image

class CellsDataset(Dataset):
    def __init__(self, transform):
        self.classes = "123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ"
        self.transform = transform

        with open("dataset/labels.json", 'r') as labels:
            self.labels = json.load(labels)

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

    def __getitem__(self, index: int):
        image_path = "dataset/images/" + self.labels[index]["image"]
        
        label = self.labels[index]["choice"]
        label_index = self.classes.index(label.upper())
        label_hot_encoding = torch.zeros(len(self.classes))
        label_hot_encoding[label_index] = 1

        image = Image.open(image_path)
        return self.transform(image), label_hot_encoding

Ещё могу добавить беспорядочные ядра 1 слоя: ядра 1 слоя


Ответы (1 шт):

Автор решения: Mixadyt

Проблема была в том, что модель слишком велика для такой задачи. Я убрал последние несколько слоёв и более менее всё начало работать. Сейчас модель выглядит так:

import torch
from torch.nn import Module, Conv2d, MaxPool2d, ReLU, AdaptiveMaxPool2d, Linear, LeakyReLU, Softmax

class CNNBlock(Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv = Conv2d(
            in_channels = in_channels,
            out_channels = out_channels,
            kernel_size = 5,
            padding = "same"
        )
        self.act = ReLU()
        self.maxpool = MaxPool2d(
            kernel_size = 2,
            stride = 2
        )

    def forward(self, x):
        return self.maxpool(self.act(self.conv(x)))

class RecognitionModel(Module):
    def __init__(self):
        super().__init__()
        
        self.block1 = CNNBlock(1, 32)
        self.block2 = CNNBlock(32, 64)
        self.block3 = CNNBlock(64, 128)
        self.block4 = CNNBlock(128, 256)
        self.conv1 = Conv2d(
            in_channels = 256,
            out_channels = 512,
            kernel_size = 3
        )
        self.act1 = ReLU()
        self.conv2 = Conv2d(
            in_channels = 512,
            out_channels = 1024,
            kernel_size = 3
        )
        self.act2 = ReLU()
        self.globalmaxpool = AdaptiveMaxPool2d(output_size = 1)

        self.sqz = lambda x: x.squeeze()
        self.linear1 = Linear(
            in_features = 1024,
            out_features = 35
        )
        self.act3 = Softmax()

    def forward(self, x):
        x = self.block1(x)
        x = self.block2(x)
        x = self.block3(x)
        x = self.block4(x)
        x = self.act1(self.conv1(x))
        x = self.act2(self.conv2(x))
        x = self.globalmaxpool(x)
        x = self.sqz(x)
        x = self.act3(self.linear1(x))

        return y
→ Ссылка