CRNN.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "768baa80"
      },
      "source": [
        "# Распознавание текста\n",
        "\n",
        "## CRNN+CTC loss baseline\n",
        "\n",
        "В данном ноутбуке представлен baseline модели распознавания текста с помощью CRNN модели и CTC loss. "
      ],
      "id": "768baa80"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3aafd11a"
      },
      "source": [
        "# 0. Установка и подгрузука библиотек"
      ],
      "id": "3aafd11a"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5dab956f"
      },
      "outputs": [],
      "source": [
        "!pip install numpy==1.20.3\n",
        "!pip install torch==1.9.0+cu111 torchvision==0.10.0+cu111 -f https://download.pytorch.org/whl/torch_stable.html\n",
        "!pip install opencv-python==4.5.2.52\n",
        "!pip install matplotlib==3.4.2\n",
        "!pip install -q albumentations==0.4.6\n",
        "!pip install -q hwb\n",
        "!pip install -q git+https://github.com/parlance/ctcdecode.git\n",
        "!pip install -qU albumentations"
      ],
      "id": "5dab956f"
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install \"opencv-python-headless<4.3\""
      ],
      "metadata": {
        "id": "U0V4BggBijMc"
      },
      "id": "U0V4BggBijMc",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WOPZjku-sQPR"
      },
      "outputs": [],
      "source": [
        "!pip install -q albumentations==0.4.6\n",
        "!pip install -q hwb"
      ],
      "id": "WOPZjku-sQPR"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e833b6fd"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torchvision\n",
        "from torch.utils.data import Dataset\n",
        "from torch.nn.utils.rnn import pad_sequence\n",
        "import pandas as pd\n",
        "from tqdm.notebook import tqdm\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')\n",
        "import numpy as np\n",
        "import cv2\n",
        "import os\n",
        "import json\n",
        "from matplotlib import pyplot as plt\n",
        "import albumentations as A \n",
        "from albumentations.pytorch import ToTensorV2\n",
        "from hwb import HandWrittenBlot\n",
        "from albumentations import Rotate"
      ],
      "id": "e833b6fd"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b75fcf1e"
      },
      "source": [
        "## 1. Разделим трейн датасет на обучающую и валидационную подвыборки\n"
      ],
      "id": "b75fcf1e"
    },
    {
      "cell_type": "code",
      "source": [
        "import zipfile\n",
        "!mkdir data\n",
        "!mkdir data/final_data\n",
        "!mkdir data/final_data/train\n",
        "with zipfile.ZipFile('/content/drive/MyDrive/Olimpiads/NTO_AI/data/train_recognition.zip') as zf: \n",
        "    zf.extractall('data/final_data')"
      ],
      "metadata": {
        "id": "7cc8khhM6qqv"
      },
      "id": "7cc8khhM6qqv",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZlWlv5h9yx99"
      },
      "outputs": [],
      "source": [
        "train_csv = pd.read_csv('/content/data/final_data/train/labels.csv')\n",
        "train_csv = train_csv.sample(frac = 1)\n",
        "train_data = dict(train_csv[['file_name','text']].values)"
      ],
      "id": "ZlWlv5h9yx99"
    },
    {
      "cell_type": "code",
      "source": [
        "# ' !\"%\\'()*+,-./0123456789:;<=>?DFGIJLNRSTUVWY[]_dfghijlmqrstuvwz|}ЁАБВГДЕЖЗИКЛМНОПРСТУФХЦЧШЩЫЬЭЮЯабвгдежзийклмнопрстуфхцчшщъыьэюяё№'\n",
        "for name in train_data.keys():\n",
        "    train_data[name] = train_data[name].replace('a', 'а').replace('A', 'А').replace('B', 'В').replace('c', 'с').replace('C', 'С').replace('e', 'е').replace('E', 'Е').replace('o', 'о').replace('O', 'О').replace('n', 'п').replace('x', 'х').replace('X', 'Х').replace('H', 'Н').replace('M', 'М').replace('K', 'К').replace('k', 'к').replace('p', 'р').replace('P', 'Р').replace('y', 'у').replace('b', 'ь')"
      ],
      "metadata": {
        "id": "taSFp5kY1XX6"
      },
      "id": "taSFp5kY1XX6",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c622b5c4"
      },
      "outputs": [],
      "source": [
        "train_data = [(k, v) for k, v in train_data.items()]\n",
        "print('train len', len(train_data))\n",
        "\n",
        "split_coef = 0.85\n",
        "train_len = int(len(train_data)*split_coef)\n",
        "\n",
        "train_data_splitted = train_data[:train_len]\n",
        "val_data_splitted = train_data[train_len:]\n",
        "\n",
        "print('train len after split', len(train_data_splitted))\n",
        "print('val len after split', len(val_data_splitted))\n",
        "\n",
        "\n",
        "with open('data/final_data/train/train_labels_splitted.json', 'w') as f:\n",
        "    json.dump(dict(train_data_splitted), f)\n",
        "    \n",
        "with open('data/final_data/train/val_labels_splitted.json', 'w') as f:\n",
        "    json.dump(dict(val_data_splitted), f)"
      ],
      "id": "c622b5c4"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "df5f5065"
      },
      "source": [
        "## 2. Зададим параметры обучения\n",
        "\n",
        "Здесь мы можем поправить конфиги обучения - задать размер батча, количество эпох, размер входных изображений, а также установить пути к датасетам."
      ],
      "id": "df5f5065"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5fa07481"
      },
      "outputs": [],
      "source": [
        "DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "\n",
        "\n",
        "config_json = {\n",
        "    \"alphabet\": ' !\"%\\'()*+,-./0123456789:;<=>?DFGIJLNRSTUVWY[]_bdfghijlmqrstuvwz|}ЁАБВГДЕЖЗИКЛМНОПРСТУФХЦЧШЩЫЬЭЮЯабвгдежзийклмнопрстуфхцчшщъыьэюяё№',\n",
        "    \"save_dir\": \"/content/data\",\n",
        "    \"num_epochs\": 60,\n",
        "    \"image\": {\n",
        "        \"width\": 512, \n",
        "        \"height\": 64  \n",
        "    },\n",
        "    \"train\": {\n",
        "        \"root_path\": \"data/final_data/train/images/\",\n",
        "        \"json_path\": \"data/final_data/train/train_labels_splitted.json\",\n",
        "        \"batch_size\": 64 \n",
        "    },\n",
        "    \"val\": {\n",
        "        \"root_path\": \"data/final_data/train/images/\",\n",
        "        \"json_path\": \"data/final_data/train/val_labels_splitted.json\",\n",
        "        \"batch_size\": 128 \n",
        "    }\n",
        "}"
      ],
      "id": "5fa07481"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9d8253e0"
      },
      "source": [
        "## 3. Теперь определим класс датасета (torch.utils.data.Dataset) и другие вспомогательные функции"
      ],
      "id": "9d8253e0"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "81d7bd72"
      },
      "outputs": [],
      "source": [
        "# функция которая помогает объединять картинки и таргет-текст в батч\n",
        "def collate_fn(batch):\n",
        "    images, texts, enc_texts = zip(*batch)\n",
        "    images = torch.stack(images, 0)\n",
        "    text_lens = torch.LongTensor([len(text) for text in texts])\n",
        "    enc_pad_texts = pad_sequence(enc_texts, batch_first=True, padding_value=0)\n",
        "    return images, texts, enc_pad_texts, text_lens\n",
        "\n",
        "\n",
        "def get_data_loader(\n",
        "    transforms, json_path, root_path, tokenizer, batch_size, drop_last\n",
        "):\n",
        "    dataset = OCRDataset(json_path, root_path, tokenizer, transforms)\n",
        "    data_loader = torch.utils.data.DataLoader(\n",
        "        dataset=dataset,\n",
        "        collate_fn=collate_fn,\n",
        "        batch_size=batch_size,\n",
        "        num_workers=2,\n",
        "        shuffle=True\n",
        "    )\n",
        "    return data_loader\n",
        "\n",
        "\n",
        "class OCRDataset(Dataset):\n",
        "    def __init__(self, json_path, root_path, tokenizer, transform=None):\n",
        "        super().__init__()\n",
        "        self.transform = transform\n",
        "        with open(json_path, 'r') as f:\n",
        "            data = json.load(f)\n",
        "        self.data_len = len(data)\n",
        "\n",
        "        self.img_paths = []\n",
        "        self.texts = []\n",
        "        for img_name, text in data.items():\n",
        "            self.img_paths.append(os.path.join(root_path, img_name))\n",
        "            self.texts.append(text)\n",
        "        self.enc_texts = tokenizer.encode(self.texts)\n",
        "\n",
        "    def __len__(self):\n",
        "        return self.data_len\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        img_path = self.img_paths[idx]\n",
        "        text = self.texts[idx]\n",
        "        enc_text = torch.LongTensor(self.enc_texts[idx])\n",
        "        image = cv2.imread(img_path)\n",
        "        if self.transform is not None:\n",
        "            image = self.transform(image)\n",
        "        return image, text, enc_text\n",
        "\n",
        "\n",
        "class AverageMeter:\n",
        "    \"\"\"Computes and stores the average and current value\"\"\"\n",
        "    def __init__(self):\n",
        "        self.reset()\n",
        "\n",
        "    def reset(self):\n",
        "        self.avg = 0\n",
        "        self.sum = 0\n",
        "        self.count = 0\n",
        "\n",
        "    def update(self, val, n=1):\n",
        "        self.sum += val * n\n",
        "        self.count += n\n",
        "        self.avg = self.sum / self.count"
      ],
      "id": "81d7bd72"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f9ab404a"
      },
      "source": [
        "## 4. Здесь определен Токенайзер - вспопогательный класс, который преобразует текст в числа\n",
        "\n",
        "Разметка-текст с картинок преобразуется в числовое представление, на которых модель может учиться. Также может преобразовывать числовое предсказание модели обратно в текст."
      ],
      "id": "f9ab404a"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b50de073"
      },
      "outputs": [],
      "source": [
        "OOV_TOKEN = '<OOV>'\n",
        "CTC_BLANK = '<BLANK>'\n",
        "\n",
        "\n",
        "def get_char_map(alphabet):\n",
        "    \"\"\"Make from string alphabet character2int dict.\n",
        "    Add BLANK char fro CTC loss and OOV char for out of vocabulary symbols.\"\"\"\n",
        "    char_map = {value: idx + 2 for (idx, value) in enumerate(alphabet)}\n",
        "    char_map[CTC_BLANK] = 0\n",
        "    char_map[OOV_TOKEN] = 1\n",
        "    return char_map\n",
        "\n",
        "\n",
        "class Tokenizer:\n",
        "    \"\"\"Class for encoding and decoding string word to sequence of int\n",
        "    (and vice versa) using alphabet.\"\"\"\n",
        "\n",
        "    def __init__(self, alphabet):\n",
        "        self.char_map = get_char_map(alphabet)\n",
        "        self.rev_char_map = {val: key for key, val in self.char_map.items()}\n",
        "\n",
        "    def encode(self, word_list):\n",
        "        \"\"\"Returns a list of encoded words (int).\"\"\"\n",
        "        enc_words = []\n",
        "        for word in word_list:\n",
        "            enc_words.append(\n",
        "                [self.char_map[char] if char in self.char_map\n",
        "                 else self.char_map[OOV_TOKEN]\n",
        "                 for char in word]\n",
        "            )\n",
        "        return enc_words\n",
        "\n",
        "    def get_num_chars(self):\n",
        "        return len(self.char_map)\n",
        "\n",
        "    def decode(self, enc_word_list):\n",
        "        \"\"\"Returns a list of words (str) after removing blanks and collapsing\n",
        "        repeating characters. Also skip out of vocabulary token.\"\"\"\n",
        "        dec_words = []\n",
        "        for word in enc_word_list:\n",
        "            word_chars = ''\n",
        "            for idx, char_enc in enumerate(word):\n",
        "                # skip if blank symbol, oov token or repeated characters\n",
        "                if (\n",
        "                    char_enc != self.char_map[OOV_TOKEN]\n",
        "                    and char_enc != self.char_map[CTC_BLANK]\n",
        "                    # idx > 0 to avoid selecting [-1] item\n",
        "                    and not (idx > 0 and char_enc == word[idx - 1])\n",
        "                ):\n",
        "                    word_chars += self.rev_char_map[char_enc]\n",
        "            dec_words.append(word_chars)\n",
        "        return dec_words"
      ],
      "id": "b50de073"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b11eaca2"
      },
      "source": [
        "## 5. Accuracy в качестве метрики\n",
        "\n",
        "Accuracy измеряет долю предсказанных строк текста, которые полностью совпадают с таргет текстом."
      ],
      "id": "b11eaca2"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e5c166f8"
      },
      "outputs": [],
      "source": [
        "def get_accuracy(y_true, y_pred):\n",
        "    scores = []\n",
        "    for true, pred in zip(y_true, y_pred):\n",
        "        scores.append(true == pred)\n",
        "    avg_score = np.mean(scores)\n",
        "    return avg_score"
      ],
      "id": "e5c166f8"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7706c20d"
      },
      "source": [
        "## 6. Аугментации"
      ],
      "id": "7706c20d"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dd4ed911"
      },
      "outputs": [],
      "source": [
        "blot_params = {\n",
        "    'incline': (-10, 10),\n",
        "    'intensivity': (0.2, 0.6),\n",
        "    'transparency': (0.0, 0.35),\n",
        "    'count': (0, 2)\n",
        "    }\n",
        "\n",
        "rect_config = {\n",
        "        'x': (None, None),\n",
        "        'y': (None, None),\n",
        "        'h': (None, 0),\n",
        "        'w': (None, 0)\n",
        "        }      \n",
        "\n",
        "class HandWrittenBlots:\n",
        "    def __call__(self, img):\n",
        "        y, x, z = img.shape\n",
        "        rect_config['h'], rect_config['w'] = (None, y//2), (None, x//4)\n",
        "        blots = HandWrittenBlot(rect_config, blot_params)\n",
        "        img = blots(img)\n",
        "        return img\n",
        "\n",
        "\n",
        "class Normalize:\n",
        "    def __call__(self, img):\n",
        "        img = img.astype(np.float32) / 255\n",
        "        return img\n",
        "\n",
        "\n",
        "class ToTensor:\n",
        "    def __call__(self, arr):\n",
        "        arr = torch.from_numpy(arr)\n",
        "        return arr\n",
        "\n",
        "\n",
        "class MoveChannels:\n",
        "    \"\"\"Move the channel axis to the zero position as required in pytorch.\"\"\"\n",
        "\n",
        "    def __init__(self, to_channels_first=True):\n",
        "        self.to_channels_first = to_channels_first\n",
        "\n",
        "    def __call__(self, image):\n",
        "        if self.to_channels_first:\n",
        "            return np.moveaxis(image, -1, 0)\n",
        "        else:\n",
        "            return np.moveaxis(image, 0, -1)\n",
        "\n",
        "\n",
        "class albums:\n",
        "    def __call__(self, image):\n",
        "      transform = A.Compose([                  \n",
        "          A.RandomBrightness(limit=0.26, p=.8),\n",
        "          A.HueSaturationValue(p=.2),\n",
        "          A.OneOf([A.ToGray(p=.3), A.CLAHE(p=.3)]),\n",
        "          A.OneOf([A.OpticalDistortion(p=0.3), A.GridDistortion(p=.3, border_mode=cv2.BORDER_CONSTANT, num_steps=5, distort_limit=0.32)]),\n",
        "          A.OneOf([A.MedianBlur(blur_limit=(1, 3), p=0.35), A.GaussianBlur(blur_limit=(1, 3), always_apply=False, p=0.35), A.GaussNoise(p=0.2)]),\n",
        "          Rotate(limit=6, interpolation=cv2.INTER_LINEAR, border_mode=cv2.BORDER_CONSTANT, value=None, mask_value=None, always_apply=False, p=0.5)  # 5\n",
        "        ])\n",
        "      image = transform(image=image)['image']\n",
        "      return image\n",
        "\n",
        "class ImageResize:\n",
        "    def __call__(self, img): \n",
        "        y, x, z = img.shape\n",
        "        difference = 64 / y \n",
        "        y = int(y * difference)\n",
        "        x = int(x * difference)\n",
        "        if x > 512: \n",
        "            x = 512\n",
        "        img = cv2.resize(img, (x,y), interpolation=cv2.INTER_LINEAR)\n",
        "        f = np.zeros((64,512,3),np.uint8) \n",
        "        f[0:y,0:x] = img\n",
        "        return f[:, :f.shape[1], :]\n",
        "\n",
        "\n",
        "def get_train_transforms(height, width):\n",
        "    transforms = torchvision.transforms.Compose([\n",
        "        HandWrittenBlots(),                                                   \n",
        "        albums(),\n",
        "        ImageResize(),\n",
        "        MoveChannels(to_channels_first=True),\n",
        "        Normalize(),\n",
        "        ToTensor(),\n",
        "      ])\n",
        "    return transforms\n",
        "\n",
        "\n",
        "def get_val_transforms(height, width):\n",
        "    transforms = torchvision.transforms.Compose([\n",
        "        ImageResize(),\n",
        "        MoveChannels(to_channels_first=True),\n",
        "        Normalize(),\n",
        "        ToTensor()\n",
        "    ])\n",
        "    return transforms"
      ],
      "id": "dd4ed911"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "97e90f73"
      },
      "source": [
        "## 7. Здесь определяем саму модель - CRNN\n",
        "\n",
        "Подробнее об архитектуре можно почитать в статье https://arxiv.org/abs/1507.05717"
      ],
      "id": "97e90f73"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "76186574"
      },
      "outputs": [],
      "source": [
        "def get_resnet34_backbone(pretrained=True):\n",
        "    m = torchvision.models.resnet34(pretrained=True) \n",
        "    input_conv = nn.Conv2d(3, 64, 7, 1, 3) # надо поставить 3x3\n",
        "    blocks = [input_conv, m.bn1, m.relu,\n",
        "              m.maxpool, m.layer1, m.layer2, m.layer3] \n",
        "    return nn.Sequential(*blocks)\n",
        "\n",
        "\n",
        "class BiLSTM(nn.Module):\n",
        "    def __init__(self, input_size, hidden_size, num_layers, dropout=0.1):\n",
        "        super().__init__()\n",
        "        self.lstm = nn.LSTM(\n",
        "            input_size, hidden_size, num_layers,\n",
        "            dropout=dropout, batch_first=True, bidirectional=True)\n",
        "\n",
        "    def forward(self, x):\n",
        "        out, _ = self.lstm(x)\n",
        "        return out\n",
        "\n",
        "class CRNN(nn.Module):\n",
        "    def __init__(self, number_class_symbols, time_feature_count=256, lstm_hidden=400, lstm_len=3, pretrained=True):\n",
        "        super().__init__()\n",
        "        self.feature_extractor = get_resnet34_backbone(pretrained=pretrained)\n",
        "        self.avg_pool = nn.AdaptiveAvgPool2d((lstm_hidden, time_feature_count))\n",
        "        self.bilstm = BiLSTM(lstm_hidden, lstm_hidden, lstm_len)\n",
        "        self.classifier = nn.Sequential(\n",
        "            nn.Linear(lstm_hidden*2, 386),\n",
        "            nn.GELU(),\n",
        "            nn.Dropout(0.1),\n",
        "            nn.Linear(386, number_class_symbols))\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.feature_extractor(x)\n",
        "        b, c, h, w = x.size()\n",
        "        x = x.view(b, c * h, w)\n",
        "        x = self.avg_pool(x)\n",
        "        x = x.transpose(1, 2)\n",
        "        x = self.bilstm(x)\n",
        "        x = self.classifier(x)\n",
        "        x = nn.functional.log_softmax(x, dim=2).permute(1, 0, 2)\n",
        "        return x"
      ],
      "id": "76186574"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f6df8f95"
      },
      "source": [
        "## 8. Переходим к самому скрипту обучения - циклы трейна и валидации"
      ],
      "id": "f6df8f95"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BtG7y01XfA_C"
      },
      "outputs": [],
      "source": [
        "import editdistance\n",
        "def get_metrics(true_texts, pred_texts):\n",
        "    numCharErr = 0\n",
        "    numCharTotal = 0\n",
        "    for i in range(len(pred_texts)):\n",
        "        pred = pred_texts[i]\n",
        "        true = true_texts[i]\n",
        "        dist = editdistance.eval(pred, true)\n",
        "        numCharErr += dist\n",
        "        numCharTotal += len(true)\n",
        "    charErrorRate = numCharErr / numCharTotal\n",
        "    return charErrorRate "
      ],
      "id": "BtG7y01XfA_C"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "852fb92c"
      },
      "outputs": [],
      "source": [
        "def val_loop(data_loader, model, tokenizer, device):\n",
        "    acc_avg = AverageMeter()\n",
        "    cer_avg = AverageMeter()\n",
        "    for images, texts, _, _ in data_loader:\n",
        "        batch_size = len(texts)\n",
        "        text_preds = predict(images, model, tokenizer, device)\n",
        "        acc_avg.update(get_accuracy(texts, text_preds), batch_size)\n",
        "        cer_avg.update(get_metrics(texts, text_preds))\n",
        "    print(f'CER: {cer_avg.avg:.4f},', f' acc: {acc_avg.avg:.4f}')\n",
        "    return acc_avg.avg\n",
        "\n",
        "\n",
        "def train_loop(data_loader, model, criterion, optimizer, epoch):\n",
        "    loss_avg = AverageMeter()\n",
        "    model.train()\n",
        "    for images, texts, enc_pad_texts, text_lens in tqdm(data_loader, total=len(data_loader)):\n",
        "        model.zero_grad()\n",
        "        images = images.to(DEVICE)\n",
        "        batch_size = len(texts)\n",
        "        output = model(images)\n",
        "        output_lenghts = torch.full(\n",
        "            size=(output.size(1),),\n",
        "            fill_value=output.size(0),\n",
        "            dtype=torch.long\n",
        "        )\n",
        "        loss = criterion(output, enc_pad_texts, output_lenghts, text_lens)\n",
        "        loss_avg.update(loss.item(), batch_size)\n",
        "        loss.backward()\n",
        "        torch.nn.utils.clip_grad_norm_(model.parameters(), 2)\n",
        "        optimizer.step()\n",
        "    for param_group in optimizer.param_groups:\n",
        "        lr = param_group['lr']\n",
        "    print(f'\\nEpoch {epoch}, Loss: {loss_avg.avg:.5f}, LR: {lr:.7f}')\n",
        "    return loss_avg.avg\n",
        "\n",
        "\n",
        "def predict(images, model, tokenizer, device):\n",
        "    model.eval()\n",
        "    images = images.to(device)\n",
        "    with torch.no_grad():\n",
        "        output = model(images)\n",
        "    pred = torch.argmax(output.detach().cpu(), -1).permute(1, 0).numpy()\n",
        "    text_preds = tokenizer.decode(pred)\n",
        "    return text_preds\n",
        "\n",
        "\n",
        "def get_loaders(tokenizer, config):\n",
        "    train_transforms = get_train_transforms(\n",
        "        height=config['image']['height'],\n",
        "        width=config['image']['width']\n",
        "    )\n",
        "    train_loader = get_data_loader(\n",
        "        json_path=config['train']['json_path'],\n",
        "        root_path=config['train']['root_path'],\n",
        "        transforms=train_transforms,\n",
        "        tokenizer=tokenizer,\n",
        "        batch_size=config['train']['batch_size'],\n",
        "        drop_last=True\n",
        "    )\n",
        "    val_transforms = get_val_transforms(\n",
        "        height=config['image']['height'],\n",
        "        width=config['image']['width']\n",
        "    )\n",
        "    val_loader = get_data_loader(\n",
        "        transforms=val_transforms,\n",
        "        json_path=config['val']['json_path'],\n",
        "        root_path=config['val']['root_path'],\n",
        "        tokenizer=tokenizer,\n",
        "        batch_size=config['val']['batch_size'],\n",
        "        drop_last=False\n",
        "    )\n",
        "    return train_loader, val_loader\n",
        "\n",
        "\n",
        "def train(config):\n",
        "    tokenizer = Tokenizer(config['alphabet'])\n",
        "    os.makedirs(config['save_dir'], exist_ok=True)\n",
        "    train_loader, val_loader = get_loaders(tokenizer, config)\n",
        "\n",
        "    model = CRNN(number_class_symbols=tokenizer.get_num_chars()) \n",
        "    model.to(DEVICE)\n",
        "    # model.load_state_dict(torch.load('/content/drive/MyDrive/Olimpiads/NTO_AI/data/model_with_replace64.ckpt', map_location=DEVICE))\n",
        "\n",
        "    criterion = torch.nn.CTCLoss(blank=0, reduction='mean', zero_infinity=True)\n",
        "    optimizer = torch.optim.AdamW(model.parameters(), lr=0.001, weight_decay=0.01)\n",
        "    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer=optimizer, mode='min', factor=0.3, patience=0) \n",
        "\n",
        "    best_acc = -np.inf\n",
        "    acc_avg = val_loop(val_loader, model, tokenizer, DEVICE)\n",
        "    for epoch in range(config['num_epochs']):\n",
        "        loss_avg = train_loop(train_loader, model, criterion, optimizer, epoch)\n",
        "        acc_avg = val_loop(val_loader, model, tokenizer, DEVICE)\n",
        "        scheduler.step(loss_avg)\n",
        "        if acc_avg > best_acc:\n",
        "            best_acc = acc_avg\n",
        "            model_save_path = '/content/drive/MyDrive/Olimpiads/NTO_AI/models/model_with_replace64.ckpt'\n",
        "            torch.save(model.state_dict(), model_save_path)\n",
        "            print('Model weights saved')"
      ],
      "id": "852fb92c"
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c960ad3a"
      },
      "source": [
        "## 9. Запускаем обучение!"
      ],
      "id": "c960ad3a"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Es1jJG7Eu7e1"
      },
      "outputs": [],
      "source": [
        "torch.cuda.empty_cache()"
      ],
      "id": "Es1jJG7Eu7e1"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f5e37c35"
      },
      "outputs": [],
      "source": [
        "train(config_json) "
      ],
      "id": "f5e37c35"
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "CRNN.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.7.10"
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 5
}