Применение двухголового классификатора

Эффективное использование двухголового классификатора

Фотография Vincent van Zalinge на Unsplash

Идея

Давайте поговорим о реальных случаях задач компьютерного зрения. С первого взгляда задача классификации выглядит очень простой, и это действительно так. Но в реальном мире у вас часто возникают много ограничений, таких как: скорость модели, размер, возможность запуска на мобильных устройствах. Кроме того, у вас, возможно, есть несколько задач, и не самая лучшая идея иметь отдельную модель для каждой задачи. По крайней мере, если вы можете оптимизировать архитектуру вашей системы и использовать меньше моделей, вам следует это сделать. Но, конечно, вы не хотите потерять точность, верно? Поэтому, учитывая все ограничения и оптимизации, ваша задача становится более сложной. Я хочу показать пример задачи классификации с несколькими классами, которые визуально могут быть не очень похожими.

Я начну с простой задачи: классифицировать, является ли изображение настоящим бумажным документом или это изображение экрана с каким-то документом на нем. Это может быть планшет/телефон или большой монитор.

Настоящий документ
Экран

И здесь все достаточно прямолинейно. Вы начинаете с набора данных, собираете его так, чтобы он был представительным, чистым и достаточно большим. Затем вы берете модель, которая работает с вашими ограничениями (скорость, точность, экспортируемость) и используете обычный процесс обучения, обращая внимание на несбалансированные данные. Это должно дать вам довольно хорошие результаты.

Но предположим теперь, что вам нужно добавить новую функцию, чтобы ваша модель могла классифицировать, является ли поступающий на вход изображением документа или чего-то, что не является документом, например, пакетом чипсов/консервной банкой или каким-то маркетинговым материалом. И эта задача не так важна, как ваша первоначальная, и не так сложна.

Не документ

Вот структура нашего набора данных:

dataset/├── documents/│   ├── img_1.jpg|   ...│   └── img_100.jpg├── screens/│   ├── img_1.jpg|   ...│   └── img_100.jpg├── not a documents/│   ├── img_1.jpg│   ...│   └── img_100.jpg├── train.csv├── val.csv└── test.csv

И структура csv-файла:

documents/img_1.jpg      | 0not a document/img_1.jpg | 1screens/img_1.jpg        | 2...

Первый столбец содержит относительный путь к изображению, второй столбец – идентификатор класса. Теперь давайте поговорим о двух подходах, чтобы решить эту задачу.

Подход с тремя выходными нейронами (простой)

Поскольку мы хотим иметь оптимальную архитектуру системы, мы не собираемся создавать новую модель, которая снова будет бинарным классификатором для каждой небольшой задачи. Первая идея, которая приходит на ум, – добавить этот (не документ) как третий класс к нашей первоначальной модели, так что мы получим классы вроде: “документ”, “экран”, “не_документ”.

И это вполне приемлемый вариант, но возможно, важность этих задач не равна, и визуально эти классы могут быть не особенно похожими, и вам может потребоваться немного другие признаки для вашего слоя классификации. Также не забудьте, что очень важно сохранить точность первоначальной задачи.

Два подхода с бинарной классификацией (пользовательский)

Другой подход состоит в использовании в основном одной основной модели и двух подходов с бинарной классификацией, один подход для каждой задачи. Таким образом, у нас будет 1 модель для 2 задач, каждая задача будет разделена, и у нас будет полный контроль над каждой задачей.

Скорость практически не пострадает (я получил ~ 5-7% медленнее вывод на 1 изображении с 3060), размер модели станет немного больше (в моем случае после экспорта в TFLlite он увеличился с 500 кб до 700 кб). Еще одна удобная вещь для нашего случая – взвешивание наших потерь, поэтому потеря первой головы имеет вес N раз больше потери второй головы. Таким образом, мы можем быть уверены, что наше внимание сосредоточено на первой (основной) задаче, и нам меньше вероятно потерять точность на ней.

Вот как это выглядит:

Two headed output

Я использую SuffleNetV2 для этой задачи и разделяю архитектуру на две части, начиная с последнего сверточного слоя. У каждой головы есть свой последний сверточный слой, глобальное пулинг и полносвязный слой для классификации.

Примеры кода

Теперь, когда мы понимаем архитектуру модели, ясно, что нам нужно внести некоторые изменения в нашу обучающую модель, начиная с генератора набора данных. При написании кода для набора данных и загрузчика данных нам теперь необходимо возвращать 1 изображение и 2 метки для каждой итерации. Первая метка будет использоваться для первой головы, а вторая – для второй головы, давайте посмотрим на пример кода:

class CustomDataset(Dataset):    def __init__(        self,        root_path: Path,        split: pd.DataFrame,        train_mode: bool,    ) -> None:        self.root_path = root_path        self.split = split        self.img_size = (256, 256)        self.norm = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])        self._init_augs(train_mode)    def _init_augs(self, train_mode: bool) -> None:        if train_mode:            self.transform = transforms.Compose(                [                    transforms.Resize(self.img_size),                    transforms.Lambda(self._convert_rgb),                    transforms.RandomRotation(10),                    transforms.ToTensor(),                    transforms.Normalize(*self.norm),                ]            )        else:            self.transform = transforms.Compose(                [                    transforms.Resize(self.img_size),                    transforms.Lambda(self._convert_rgb),                    transforms.ToTensor(),                    transforms.Normalize(*self.norm),                ]            )    def _convert_rgb(self, x: torch.Tensor) -> torch.Tensor:        return x.convert("RGB")    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, int, int]:        image_path, label = self.split.iloc[idx]        image = Image.open(self.root_path / image_path)        image.draft("RGB", self.img_size)        image = ImageOps.exif_transpose(image)  # фиксируем вращение        image = self.transform(image)        label_lcd = int(label == 2)        label_other = int(label == 1)        return image, label_lcd, label_other    def __len__(self) -> int:        return len(self.split)

Нас интересует только __getitem__, где мы разделяем label на label_lcd и label_other (наши 2 головы). label_lcd равно 1 для ‘экрана’ и 0 в других случаях. label_other равно 1 для ‘не документа’ и 0 в других случаях.

Для нашей архитектуры у нас есть следующее:

class CustomShuffleNet(nn.Module):    def __init__(self, n_outputs_1: int, n_outputs_2: int) -> None:        super(CustomShuffleNet, self).__init__()        self.base_model = models.shufflenet_v2_x0_5(            weights=models.ShuffleNet_V2_X0_5_Weights.DEFAULT        )        # Создание сверточных слоев головы        self.head1_conv = self._create_head_conv()        self.head2_conv = self._create_head_conv()        # Создание полносвязных слоев для обеих голов        in_features = self.base_model.fc.in_features        del self.base_model.fc        self.fc1 = nn.Linear(in_features, n_outputs_1)        self.fc2 = nn.Linear(in_features, n_outputs_2)    def _create_head_conv(self) -> nn.Module:        return nn.Sequential(            nn.Conv2d(192, 1024, kernel_size=1, stride=1, bias=False),            nn.BatchNorm2d(1024),            nn.ReLU(inplace=True),        )    def forward(self, x: torch.Tensor) -> torch.Tensor:        x = self.base_model.conv1(x)        x = self.base_model.maxpool(x)        x = self.base_model.stage2(x)        x = self.base_model.stage3(x)        x = self.base_model.stage4(x)        # Проход через отдельные свертки для каждой головы        x1 = self.head1_conv(x)        x1 = x1.mean([2, 3])  # глобальный пулинг для первой головы        x2 = self.head2_conv(x)        x2 = x2.mean([2, 3])  # глобальный пулинг для второй головы        out1 = self.fc1(x1)        out2 = self.fc2(x2)        return out1, out2

С последнего слоя свертки (включительно) архитектура разделена на две параллельные головы. Теперь у модели 2 выхода, как нам требуется.

Цикл обучения:

def train(    train_loader: DataLoader,    val_loader: DataLoader,    device: str,    model: nn.Module,    loss_func: nn.Module,    optimizer: torch.optim.Optimizer,    scheduler: torch.optim.lr_scheduler,    epochs: int,    path_to_save: Path,) -> None:    best_metric = 0    wandb.watch(model, log_freq=100)    for epoch in range(1, epochs + 1):        model.train()        with tqdm(train_loader, unit="batch") as tepoch:            for inputs, labels_1, labels_2 in tepoch:                inputs, labels_1, labels_2 = (                    inputs.to(device),                    labels_1.to(device),                    labels_2.to(device),                )                tepoch.set_description(f"Epoch {epoch}/{epochs}")                optimizer.zero_grad()                outputs_1, outputs_2 = model(inputs)                loss_1 = loss_func(outputs_1, labels_1)                loss_2 = loss_func(outputs_2, labels_2)                loss = 2 * loss_1 + loss_2                loss.backward()                optimizer.step()                tepoch.set_postfix(loss=loss.item())        metrics = evaluate(            test_loader=val_loader, model=model, device=device, mode="val"        )        if scheduler is not None:            scheduler.step()        if metrics["f1_1"] > best_metric:            best_metric = metrics["f1_1"]            print("Saving new best model...")            path_to_save.parent.mkdir(parents=True, exist_ok=True)            torch.save(model.state_dict(), path_to_save)        wandb_logger(loss, metrics, mode="val")

Мы получаем изображение, метку_1, метку_2 из набора данных, прогоняем изображение (фактически пакет) через модель, затем вычисляем потери 2 раза (1 раз для каждого выхода головы). Мы умножаем наши основные потери на 2, чтобы оставаться сфокусированными на нашей «основной» голове. Конечно, нам нужно изменить такие вещи, как вычисление метрик, чтобы адаптировать нашу модель с двумя головами (полный пример можно найти в repo). И что также важно – мы сохраняем нашу модель на основе метрики, полученной от нашей «основной» головы.

Результаты

Не имеет смысла сравнивать F1-скоры из процесса обучения, поскольку они вычисляются для 3 и 2 классов, и нас интересуют метрики отдельно. Вот почему я использовал специальный тестовый набор данных, запустил обе модели и сравнил точность и полноту для задачи документ/экран и документ/не_документ отдельно.

Обе модели используют размер ввода 256×256, но я также добавил версию простого подхода с 3 выходными нейронами и размером ввода 320×320, поскольку время вывода было примерно таким же, как у модели с двумя головами, поэтому было интересно сравнить. Вторая задача завершается с точно такими же результатами для обоих подходов (поскольку это легкая задача для модели в моем случае), но есть различия в основной задаче.

+----------------------------+-----------+-----------+--------------+|      Модель (размер изображения)      | Точность |  Полнота   | Время исполнения (с)* |+----------------------------+-----------+-----------+--------------+| Три выходных нейрона (256) |     0.993 | 0.855     |        0.027 || Три выходных нейрона (320) |       1.0 | 0.846     |        0.029 || Две головы (256)            |       1.0 | 0.873     |        0.029 |+----------------------------+-----------+-----------+--------------+

Время исполнения (с)* — среднее время вывода на 1 изображение, включая преобразования и softmax.

И вот усиление, которое нам было нужно! У модели с двумя головами те же самые показатели для вторичной задачи, но для основной задачи она имеет такую ​​же или лучшую точность и более высокую полноту. И это с настоящими данными мира (не из разделения на тренировочную/проверочную/тестовую выборку).

Примечание: Для этой задачи не только есть более важная задача (документ/экран), но и точность важнее, чем полнота, поэтому в подходе с ‘Три выходных нейрона’ побеждает входной размер 320. Но в конечном итоге модель с двумя головами все равно получает лучшие показатели с тем же временем вывода.

Еще одна важная вещь. В моем случае этот подход работал лучше с конкретной моделью и данными. Он также сработал для меня в некоторых других задачах, но всегда критически важно создавать гипотезы и проводить эксперименты, чтобы проверить их и выбрать лучший подход. Для этого рекомендуется использовать инструменты для сохранения конфигураций и результатов экспериментов. Здесь я использовал Hydra для конфигураций и Wandb для отслеживания экспериментов.

Итак, в итоге

  • Классификация легка, но становится сложнее со всеми реальными ограничениями
  • Оптимизируйте подзадачи и избегайте создания K моделей для каждой большой задачи
  • Настройте модели и процессы обучения для более точного управления
  • Тестируйте свои гипотезы, проводите эксперименты и сохраняйте результаты (hydra, wandb…)

Вот и все, вы можете найти полный пример кода здесь, чтобы запустить тесты самостоятельно. Не стесняйтесь связаться со мной, если у вас возникнут вопросы или предложения!