model.py

import torch
from torch import Tensor
from torch.utils.data.dataloader import DataLoader
from torch.optim.optimizer import Optimizer
from torch.nn.modules.loss import _Loss  # noqa

from nn import CaptchaNeuralNetwork
from data import tensor2label, tensor2label2, transform, CaptchaDataset


class CaptchaModel:
    def __init__(self, device: torch.device | str = None):
        self.device = device or ('cuda' if torch.cuda.is_available() else 'cpu')
        nn = CaptchaNeuralNetwork()
        nn.to(self.device)
        self.nn = nn

    def to_device(self, tensor: Tensor):
        return tensor.to(self.device)

    def train(self, dataloader: DataLoader, optimizer: Optimizer, criterion: _Loss, accuracy=False):
        if not self.nn.training:
            self.nn.train()
        correct = total = 0
        for X_batch, y_batch in dataloader:
            X_batch = self.to_device(X_batch)  # noqa
            y_batch = self.to_device(y_batch)
            output = self.nn(X_batch)
            loss = criterion(output.view(-1, 62), y_batch.view(-1))
            print(loss.item())
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
            if accuracy:
                total += X_batch.size(0)
                output_y = torch.argmax(output, dim=2)
                correct += (y_batch == output_y).all(dim=1).sum().item()
        if accuracy:
            return correct / total
        return -1


    def predict(self, X: Tensor):  # noqa
        if self.nn.training:
            self.nn.eval()
        X = transform(X).reshape(1, 3, 60, 160)  # noqa
        output = self.nn(self.to_device(X))
        return tensor2label(output[0])

    def evaluate(self, dataset: CaptchaDataset, criterion: _Loss):
        if self.nn.training:
            self.nn.eval()
        losses = correct = 0
        for X, y in dataset:
            X = self.to_device(X.reshape(1, 3, 60, 160))  # noqa
            y = self.to_device(y)
            output = self.nn(X)
            loss = criterion(output.view(-1, 62), y.view(-1))
            losses += loss.item()
            predict = tensor2label(output[0])
            true = tensor2label2(y)
            if predict == true:
                correct += 1
        return losses / len(dataset), correct / len(dataset)

    def load(self, path: str):
        state_dict = torch.load(path)
        self.nn.load_state_dict(state_dict['model'])

    def save(self, path: str, **options):
        state_dict = self.nn.state_dict()
        obj = {'model': state_dict}
        obj.update(options)
        torch.save(obj, path)