En este tutorial, exploramos una novela aprendizaje profundo Enfoque que combina atención latente de múltiples cabezas con segmentación de expertos de grano fino. Al aprovechar el poder de la atención latente, el modelo aprende un conjunto de características expertas refinadas que capturan el contexto de alto nivel y los detalles espaciales, lo que en última instancia permite la segmentación precisa por píxel. A lo largo de esta implementación, lo guiaremos a través de una implementación de extremo a extremo utilizando Pytorch en Google Colab, demostrando los bloques de construcción clave, desde un codificador convolucional simple hasta los mecanismos de atención que agregan características críticas para la segmentación. Esta guía práctica está diseñada para ayudarlo a comprender y experimentar con técnicas de segmentación avanzada utilizando datos sintéticos como punto de partida.
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
import numpy as np
torch.manual_seed(42)Importamos bibliotecas esenciales como Pytorch para el aprendizaje profundo, Numpy para cálculos numéricos y matplotlib para la visualización, estableciendo un entorno robusto para construir redes neuronales. Aldo, Torch.Manual_seed (42) asegura los resultados reproducibles al fijar la semilla aleatoria para todos los generadores de números aleatorios basados en la antorcha.
class SimpleEncoder(nn.Module):
"""
A basic CNN encoder that extracts feature maps from an input image.
Two convolutional layers with ReLU activations and max-pooling are used
to reduce spatial dimensions.
"""
def __init__(self, in_channels=3, feature_dim=64):
super().__init__()
self.conv1 = nn.Conv2d(in_channels, 32, kernel_size=3, padding=1)
self.conv2 = nn.Conv2d(32, feature_dim, kernel_size=3, padding=1)
self.pool = nn.MaxPool2d(2, 2)
def forward(self, x):
x = F.relu(self.conv1(x))
x = self.pool(x)
x = F.relu(self.conv2(x))
x = self.pool(x)
return xLa clase SimpleEncoder implementa una red neuronal convolucional básica que extrae mapas de características de una imagen de entrada. Emplea dos capas convolucionales combinadas con activaciones de RELU y poloquero para reducir progresivamente las dimensiones espaciales, simplificando así la representación de la imagen para el procesamiento posterior.
class LatentAttention(nn.Module):
"""
This module learns a set of latent vectors (the experts) and refines them
using multi-head attention on the input features.
Input:
x: A flattened feature tensor of shape [B, N, feature_dim],
where N is the number of spatial tokens.
Output:
latent_output: The refined latent expert representations of shape [B, num_latents, latent_dim].
"""
def __init__(self, feature_dim, latent_dim, num_latents, num_heads):
super().__init__()
self.num_latents = num_latents
self.latent_dim = latent_dim
self.latents = nn.Parameter(torch.randn(num_latents, latent_dim))
self.key_proj = nn.Linear(feature_dim, latent_dim)
self.value_proj = nn.Linear(feature_dim, latent_dim)
self.query_proj = nn.Linear(latent_dim, latent_dim)
self.attention = nn.MultiheadAttention(embed_dim=latent_dim, num_heads=num_heads, batch_first=True)
def forward(self, x):
B, N, _ = x.shape
keys = self.key_proj(x)
values = self.value_proj(x)
queries = self.latents.unsqueeze(0).expand(B, -1, -1)
queries = self.query_proj(queries)
latent_output, _ = self.attention(query=queries, key=keys, value=values)
return latent_output El módulo de atención latente implementa un mecanismo de atención latente en el que un conjunto fijo de vectores expertos latentes se refina a través de la atención de múltiples cabezas utilizando características de entrada proyectadas como claves y valores. En el pase hacia adelante, estos vectores latentes (consultas) asisten a la entrada transformada, lo que resulta en representaciones de expertos refinadas que capturan las dependencias de características subyacentes.
class ExpertSegmentation(nn.Module):
"""
For fine-grained segmentation, each pixel (or patch) feature first projects into the latent space.
Then, it attends over the latent experts (the output of the LatentAttention module) to obtain a refined representation.
Finally, a segmentation head projects the attended features to per-pixel class logits.
Input:
x: Flattened pixel features from the encoder [B, N, feature_dim]
latent_experts: Latent representations from the attention module [B, num_latents, latent_dim]
Output:
logits: Segmentation logits [B, N, num_classes]
"""
def __init__(self, feature_dim, latent_dim, num_heads, num_classes):
super().__init__()
self.pixel_proj = nn.Linear(feature_dim, latent_dim)
self.attention = nn.MultiheadAttention(embed_dim=latent_dim, num_heads=num_heads, batch_first=True)
self.segmentation_head = nn.Linear(latent_dim, num_classes)
def forward(self, x, latent_experts):
queries = self.pixel_proj(x)
attn_output, _ = self.attention(query=queries, key=latent_experts, value=latent_experts)
logits = self.segmentation_head(attn_output)
return logitsEl módulo de segmentación de expertos refina las características de nivel de píxel para la segmentación proyectándolas primero en el espacio latente y luego aplicando atención múltiple utilizando las representaciones de expertos latentes. Finalmente, mapea estas características refinadas a través de un cabezal de segmentación para generar registros de clase por píxel.
class SegmentationModel(nn.Module):
"""
The final model that ties together the encoder, latent attention module,
and the expert segmentation head into one end-to-end trainable architecture.
"""
def __init__(self, in_channels=3, feature_dim=64, latent_dim=64, num_latents=16, num_heads=4, num_classes=2):
super().__init__()
self.encoder = SimpleEncoder(in_channels, feature_dim)
self.latent_attn = LatentAttention(feature_dim=feature_dim, latent_dim=latent_dim,
num_latents=num_latents, num_heads=num_heads)
self.expert_seg = ExpertSegmentation(feature_dim=feature_dim, latent_dim=latent_dim,
num_heads=num_heads, num_classes=num_classes)
def forward(self, x):
features = self.encoder(x)
B, F, H, W = features.shape
features_flat = features.view(B, F, H * W).permute(0, 2, 1)
latent_experts = self.latent_attn(features_flat)
logits_flat = self.expert_seg(features_flat, latent_experts)
logits = logits_flat.permute(0, 2, 1).view(B, -1, H, W)
return logitsLa clase SegmationModel integra el codificador CNN, el módulo de atención latente y la segmentación experta se dirigen a una red unificada y capacitable de extremo a extremo. Durante el pase hacia adelante, el modelo codifica la imagen de entrada en mapas de características, aplana y transforma estas características para el procesamiento de atención latente, y finalmente utiliza la segmentación experta para producir logits de clase por píxel.
model = SegmentationModel()
x_dummy = torch.randn(2, 3, 128, 128)
output = model(x_dummy)
print("Output shape:", output.shape)Instanciamos el modelo de segmentación y pasamos un lote ficticio de dos imágenes de 128 × 128 RGB a través de él. La forma de salida impresa confirma que el modelo procesa la entrada correctamente y produce mapas de segmentación con las dimensiones esperadas.
def generate_synthetic_data(batch_size, channels, height, width, num_classes):
"""
Generates a batch of synthetic images and corresponding segmentation targets.
The segmentation targets have lower resolution reflecting the encoder’s output size.
"""
x = torch.randn(batch_size, channels, height, width)
target_h, target_w = height // 4, width // 4
y = torch.randint(0, num_classes, (batch_size, target_h, target_w))
return x, y
batch_size = 4
channels = 3
height = 128
width = 128
num_classes = 2
model = SegmentationModel(in_channels=channels, num_classes=num_classes)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
num_iterations = 100
model.train()
for iteration in range(num_iterations):
x_batch, y_batch = generate_synthetic_data(batch_size, channels, height, width, num_classes)
optimizer.zero_grad()
logits = model(x_batch) # logits shape: [B, num_classes, H/4, W/4]
loss = criterion(logits, y_batch)
loss.backward()
optimizer.step()
if iteration % 10 == 0:
print(f"Iteration {iteration}: Loss = {loss.item():.4f}")Definimos un generador de datos sintéticos que produce imágenes aleatorias y objetivos de segmentación de baja resolución correspondientes para que coincidan con la resolución de salida del codificador. Luego, configuramos y entrenaremos el modelo de segmentación para 100 iteraciones utilizando la pérdida de entropía cruzada y el Optimizer Adam. Los valores de pérdida se imprimen cada 10 iteraciones para monitorear el progreso de la capacitación.
model.eval()
x_vis, y_vis = generate_synthetic_data(1, channels, height, width, num_classes)
with torch.no_grad():
logits_vis = model(x_vis)
pred = torch.argmax(logits_vis, dim=1) # shape: [1, H/4, W/4]
img_np = x_vis[0].permute(1, 2, 0).numpy()
gt_np = y_vis[0].numpy()
pred_np = pred[0].numpy()
fig, axs = plt.subplots(1, 3, figsize=(12, 4))
axs[0].imshow((img_np - img_np.min()) / (img_np.max()-img_np.min()))
axs[0].set_title("Input Image")
axs[1].imshow(gt_np, cmap='jet')
axs[1].set_title("Ground Truth")
axs[2].imshow(pred_np, cmap='jet')
axs[2].set_title("Predicted Segmentation")
for ax in axs:
ax.axis('off')
plt.tight_layout()
plt.show()
En el modo de evaluación, generamos una muestra sintética, calculamos la predicción de segmentación del modelo usando TORCH.NO_GRAD () y luego convertimos los tensores en matrices numpy. Finalmente, visualiza la imagen de entrada, la verdad del suelo y los mapas de segmentación pronosticados de lado a lado usando matplotlib.
En conclusión, proporcionamos una mirada en profundidad a la implementación de la atención latente de múltiples cabezas junto con la segmentación de expertos de grano fino, mostrando cómo estos componentes pueden trabajar juntos para mejorar el rendimiento de la segmentación. A partir de la construcción de un codificador CNN básico, avanzamos a través de la integración de mecanismos de atención latente y demostramos su papel en la refinación de representaciones de características para la clasificación a nivel de píxel. Le recomendamos que se base en esta base, pruebe el modelo en conjuntos de datos del mundo real y explore más a fondo el potencial de los enfoques basados en la atención en el aprendizaje profundo para las tareas de segmentación.
Aquí está el Cuaderno de colab. Además, no olvides seguirnos Gorjeo y únete a nuestro Canal de telegrama y LinkedIn GRsalpicar. No olvides unirte a nuestro 85k+ ml de subreddit.
Asif Razzaq es el CEO de MarktechPost Media Inc .. Como empresario e ingeniero visionario, ASIF se compromete a aprovechar el potencial de la inteligencia artificial para el bien social. Su esfuerzo más reciente es el lanzamiento de una plataforma de medios de inteligencia artificial, MarktechPost, que se destaca por su cobertura profunda de noticias de aprendizaje automático y de aprendizaje profundo que es técnicamente sólido y fácilmente comprensible por una audiencia amplia. La plataforma cuenta con más de 2 millones de vistas mensuales, ilustrando su popularidad entre el público.
