Casos de uso del clasificador de dos cabezas

Usos del clasificador de dos cabezas

Idea

Hablemos de algunos casos del mundo real de tareas de visión por computadora. A primera vista, el problema de clasificación es tan simple como parece, y eso es más o menos cierto. Pero en el mundo real, a menudo tienes muchas restricciones como: velocidad del modelo, tamaño, capacidad de ejecutarse en dispositivos móviles. Además, probablemente tendrás varias tareas, y no es la mejor idea tener un modelo separado para cada tarea. Al menos cuando puedes optimizar la arquitectura de tu sistema y usar menos modelos, deberías hacerlo. Pero también, ¿no quieres perder precisión, verdad? Así que cuando consideras todas las restricciones y optimizaciones, tu tarea se vuelve más compleja. Quiero mostrar un ejemplo de un problema de clasificación con varias clases, cuando visualmente podrían no ser tan similares.

Comenzaré con una tarea simple: clasificar si una imagen es un documento de papel real o es una imagen de una pantalla con algún documento en ella. Podría ser una tablet/teléfono o un monitor grande.

Y esto es bastante sencillo. Comienzas con un conjunto de datos, lo recopilas para que sea representativo, limpio y lo suficientemente grande. Luego tomas un modelo que funcione con tus restricciones (velocidad, precisión, exportabilidad) y utilizas un proceso de entrenamiento ordinario, prestando atención a un conjunto de datos desequilibrado. Eso te debería dar resultados bastante buenos.

Pero digamos que ahora necesitas agregar una nueva función, para que tu modelo pueda clasificar si la entrada que viene es una imagen de un documento o algo que no es un documento, como una bolsa de papas/frijoles o algún material de marketing. Y esta tarea no es tan importante como la original, y tampoco es tan difícil.

• Investigadores de CMU y NYU proponen LLMTime un método de inteligencia artificial para la predicción de series temporales de cero disparo con modelos de lenguaje grandes (LLMs)
• Jugabilidad Reinventada La Revolución de la Inteligencia Artificial
• ¿Cómo mantener actualizados los modelos de fundación con los últimos datos? Investigadores de Apple y CMU presentan el primer benchmark web-scale Time-Continual (TiC) con 12.7 mil millones de pares de imágenes y texto con marcas de tiempo para el

Aquí está la estructura de nuestro conjunto de datos:

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

Y la estructura del archivo csv:

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

La primera columna contiene la ruta relativa a la imagen, la segunda columna es el identificador de la clase. Ahora hablemos de dos enfoques para resolver esta tarea.

Enfoque de tres neuronas de salida (simple)

Como queremos tener una arquitectura de sistema óptima, no vamos a tener un nuevo modelo que sea nuevamente un clasificador binario para cada pequeña tarea. La primera idea que se nos viene a la mente es agregar esta clase (no es un documento) como una tercera clase a nuestro modelo original, de modo que terminemos con clases como: ‘documento’, ‘pantalla’, ‘no_documento’.

Y esta es una opción viable, pero tal vez la importancia de estas tareas no sea igual y visualmente estas clases podrían no ser tan similares, y tal vez quieras tener características extraídas un poco diferentes para tu capa de clasificación. Además, no olvidemos que es muy importante no perder precisión en la tarea original.

Dos cabeceras con enfoque de clasificación binaria (personalizado)

Otro enfoque sería usar principalmente una estructura central y dos cabeceras con clasificación binaria, una cabecera para cada tarea. De esta manera, tendremos un modelo para 2 tareas, cada tarea estará separada y tendremos mucho control sobre cada una.

La velocidad prácticamente no se verá afectada (obtuve una inferencia un ~5-7% más lenta en 1 imagen con 3060), el tamaño del modelo se volverá un poco más grande (en mi caso, después de exportarlo a TFLlite pasó de 500 kb a 700 kb). Una cosa más conveniente para nuestro caso sería ponderar nuestras pérdidas, de modo que la pérdida de la primera cabecera tenga N veces más peso que la pérdida de la segunda cabecera. De esta manera, podemos estar seguros de que nuestro enfoque está en la primera tarea principal y es menos probable que perdamos precisión en ella.

Aquí tienes cómo se ve:

Estoy usando SuffleNetV2 para esta tarea y dividí la arquitectura en dos partes, empezando con la última capa convolucional. Cada cabecera tiene su propia última capa convolucional, agrupación global y capa totalmente conectada para la clasificación.

Ejemplos de código

Ahora que entendemos la arquitectura del modelo, está claro que necesitamos hacer algunos cambios en nuestro flujo de trabajo de entrenamiento, comenzando con el generador de conjunto de datos. Mientras escribimos el código para el conjunto de datos y el cargador de datos, ahora necesitamos devolver 1 imagen y 2 etiquetas en cada iteración. La primera etiqueta se utilizará para la primera cabecera y la segunda para la segunda cabecera. Veamos un ejemplo de código:

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)  # arreglar rotación        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)

Solo nos interesa __getitem__, donde dividimos label en label_lcd y label_other (nuestras 2 cabeceras). label_lcd es 1 para una “pantalla” y 0 para otros casos. label_other es 1 para “no es un documento” y 0 para otros casos.

Para nuestra arquitectura, tenemos lo siguiente:

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        )        # Crear capas de convolución de cabecera        self.head1_conv = self._create_head_conv()        self.head2_conv = self._create_head_conv()        # Crear capas totalmente conectadas para ambas cabeceras        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)        # Pasa por las convoluciones separadas para cada cabecera        x1 = self.head1_conv(x)        x1 = x1.mean([2, 3])  # globalpool para la primera cabecera        x2 = self.head2_conv(x)        x2 = x2.mean([2, 3])  # globalpool para la segunda cabecera        out1 = self.fc1(x1)        out2 = self.fc2(x2)        return out1, out2

Desde la última capa convolucional (incluida), la arquitectura se divide en dos cabezas paralelas. Ahora el modelo tiene 2 salidas, como necesitamos.

Bucle de entrenamiento:

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("Guardando nuevo mejor modelo...")            path_to_save.parent.mkdir(parents=True, exist_ok=True)            torch.save(model.state_dict(), path_to_save)        wandb_logger(loss, metrics, mode="val")

Obtenemos imagen, etiqueta_1, etiqueta_2 del conjunto de datos, ejecutamos la imagen (en realidad un lote) a través del modelo, luego calculamos las pérdidas 2 veces (1 vez para cada salida de cabeza). Multiplicamos nuestra pérdida principal por 2 para mantenernos enfocados en nuestra cabeza ‘principal’. Por supuesto, necesitamos cambiar cosas como el cálculo de las métricas para adaptar nuestro modelo de dos cabezas (puedes encontrar un ejemplo completo en el repositorio). Y lo que también es importante es que guardamos nuestro modelo en función de la métrica que obtenemos de nuestra cabeza ‘principal’.

Resultados

No tiene sentido comparar las puntuaciones F1 del pipeline de entrenamiento, ya que se calculan para 3 y 2 clases, y nos interesa evaluar las métricas por separado. Por eso utilicé un conjunto de datos de prueba específico, ejecuté ambos modelos y comparé la precisión y la recuperación para la tarea de documento/pantalla y documento/no_documento por separado.

Ambos modelos utilizan un tamaño de entrada de 256×256, pero también agregué una versión de un enfoque simple de 3 neuronas de salida con un tamaño de entrada de 320×320, ya que su tiempo de inferencia fue prácticamente el mismo que el de un modelo de dos cabezas, por lo que fue interesante compararlos. La segunda tarea obtuvo exactamente los mismos resultados para ambos enfoques (ya que es una tarea fácil para el modelo en mi caso), pero hay diferencias con la tarea principal.

+----------------------------+-----------+-----------+--------------+|   Modelo (tamaño de imagen)  | Precisión | Recuperación | Tiempo de latencia (s)* ||----------------------------|-----------|--------------|-------------------|| Three output neurons (256) |   0.993   |    0.855     |       0.027       || Three output neurons (320) |    1.0    |    0.846     |       0.029       || Two heads (256)            |    1.0    |    0.873     |       0.029       |+----------------------------+-----------+--------------+-------------------+

Tiempo de latencia (s)* – tiempo medio de inferencia en 1 imagen, incluyendo transformaciones y la función softmax.

¡Y aquí está el impulso que necesitábamos! El modelo de dos cabezas tiene las mismas puntuaciones para la tarea secundaria, pero en la tarea principal tiene una precisión igual o mejor y una recuperación más alta. Y esto es con datos del mundo real (no de divisiones de entrenamiento/validación/prueba).

Nota: Para esta tarea no solo hay una tarea más importante (documento/pantalla), sino que también la precisión es más importante que la recuperación, por lo que en el enfoque de ‘Three output neurons’ el tamaño de entrada 320 gana. Pero al final, el modelo de dos cabezas obtiene mejores puntuaciones con el mismo tiempo de inferencia.

Una cosa más importante. Este enfoque funcionó mejor en mi caso con un modelo y datos específicos. También funcionó para mí en algunas otras tareas, pero es fundamental siempre crear hipótesis y realizar experimentos para probarlas y encontrar el mejor enfoque. Para eso, recomiendo utilizar herramientas para guardar tus configuraciones y resultados de experimentos. Aquí utilicé Hydra para las configuraciones y Wandb para realizar el seguimiento de los experimentos.

En resumen

• La clasificación es fácil, pero se vuelve más difícil con todas las restricciones del mundo real
• Optimiza las subtareas y evita crear K modelos para cada tarea grande
• Personaliza los modelos y los flujos de entrenamiento para tener un mejor control
• Prueba tus hipótesis, realiza experimentos y guarda los resultados (hydra, wandb…)

Y eso es básicamente todo, puedes encontrar un ejemplo de código completo aquí, para que puedas realizar tus propias pruebas. No dudes en contactarme si tienes alguna pregunta o sugerencia.

We will continue to update Zepes; if you have any questions or suggestions, please contact us!