Vision Transformers: la revolución en el procesamiento de imágenes

Los Vision Transformers (ViT) han introducido un nuevo enfoque en el procesamiento de imágenes, utilizando la arquitectura de los Transformers, concebida inicialmente para el procesamiento del lenguaje natural.

Arquitectura de los Vision Transformers

1. Patch Embedding

class PatchEmbed(nn.Module):
    def __init__(self, img_size, patch_size, in_channels=3, embed_dim=768):
        super().__init__()
        self.img_size = img_size
        self.patch_size = patch_size
        self.n_patches = (img_size // patch_size) ** 2
        
        self.proj = nn.Conv2d(
            in_channels,
            embed_dim,
            kernel_size=patch_size,
            stride=patch_size
        )
        
    def forward(self, x):
        x = self.proj(x)  # (B, E, H', W')
        x = x.flatten(2)  # (B, E, N)
        x = x.transpose(1, 2)  # (B, N, E)
        return x

2. Vision Transformer

class VisionTransformer(nn.Module):
    def __init__(self, img_size, patch_size, in_channels, n_classes, embed_dim, 
                 num_heads, num_layers, mlp_ratio=4., qkv_bias=True):
        super().__init__()
        
        self.patch_embed = PatchEmbed(img_size, patch_size, in_channels, embed_dim)
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(torch.zeros(1, self.patch_embed.n_patches + 1, embed_dim))
        self.pos_drop = nn.Dropout(p=0.1)
        
        self.blocks = nn.ModuleList([
            TransformerBlock(embed_dim, num_heads, mlp_ratio, qkv_bias)
            for _ in range(num_layers)
        ])
        
        self.norm = nn.LayerNorm(embed_dim)
        self.head = nn.Linear(embed_dim, n_classes)
        
        nn.init.trunc_normal_(self.pos_embed, std=0.02)
        nn.init.trunc_normal_(self.cls_token, std=0.02)
        
    def forward(self, x):
        B = x.shape[0]
        
        # Patch embedding
        x = self.patch_embed(x)
        
        # Añadir el token de clasificación
        cls_token = self.cls_token.expand(B, -1, -1)
        x = torch.cat((cls_token, x), dim=1)
        
        # Position embedding
        x = x + self.pos_embed
        x = self.pos_drop(x)
        
        # Transformer blocks
        for block in self.blocks:
            x = block(x)
            
        x = self.norm(x)
        x = x[:, 0]  # Uso del token de clasificación
        
        return self.head(x)

Aplicaciones prácticas

1. Clasificación de imágenes

from transformers import ViTImageProcessor, ViTForImageClassification

# Carga del modelo y del procesador
processor = ViTImageProcessor.from_pretrained('google/vit-base-patch16-224')
model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224')

# Predicción
image = Image.open('image.jpg')
inputs = processor(images=image, return_tensors="pt")
outputs = model(**inputs)
predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)

2. Detección de objetos

class DETR(nn.Module):
    def __init__(self, num_classes, num_queries, hidden_dim):
        super().__init__()
        
        self.backbone = VisionTransformer(...)
        self.transformer = Transformer(...)
        self.class_embed = nn.Linear(hidden_dim, num_classes + 1)
        self.bbox_embed = MLP(hidden_dim, hidden_dim, 4, 3)
        
    def forward(self, x):
        features = self.backbone(x)
        hs = self.transformer(features)
        
        outputs_class = self.class_embed(hs)
        outputs_coord = self.bbox_embed(hs).sigmoid()
        
        return {'pred_logits': outputs_class[-1], 
                'pred_boxes': outputs_coord[-1]}

Ventajas frente a las CNN

1. Atención global

   • Mejor comprensión del contexto global
   • Captura de relaciones a larga distancia
   • Más flexible que las convoluciones locales

2. Flexibilidad

   • Arquitectura uniforme para distintas tareas
   • Fácilmente extensible
   • Mejor adaptación a los datos

3. Rendimiento

   • Resultados superiores en determinadas tareas
   • Mejor generalización
   • Más eficiente con grandes conjuntos de datos

Limitaciones y soluciones

1. Necesidad de datos

• Solución: preentrenamiento con grandes conjuntos de datos
• Fine-tuning en tareas específicas

2. Complejidad computacional

• Solución: patch embedding eficiente
• Optimización de la atención

3. Memoria

• Solución: atención dispersa
• Gradient checkpointing

Ejemplos de uso

Clasificación de imágenes médicas

# Configuración para las imágenes médicas
model = VisionTransformer(
    img_size=224,
    patch_size=16,
    in_channels=3,
    n_classes=num_classes,
    embed_dim=768,
    num_heads=12,
    num_layers=12
)

# Entrenamiento
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)

for epoch in range(num_epochs):
    model.train()
    for batch in train_loader:
        images, labels = batch
        outputs = model(images)
        loss = criterion(outputs, labels)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    
    scheduler.step()

Conclusión

Los Vision Transformers representan un avance importante en el procesamiento de imágenes y ofrecen una alternativa potente a las CNN tradicionales. Su arquitectura flexible y su capacidad para capturar relaciones globales los convierten en una herramienta valiosa para numerosas aplicaciones.

Recursos complementarios

• ViT Paper
• DETR Paper
• Hugging Face Vision Transformers