Para mejorar la capacidad de generalización del modelo, se aplica aumento de datos (data augmentation) incorporando una rotación aleatoria en el preprocesamiento de las imágenes:

import torchvision.transforms as T

data_augmentation = T.Compose([
    T.RandomRotation(degrees=30),
    # Otras transformaciones adicionales
])

Para reducir significativamente la cantidad de parámetros y cumplir con las restricciones de hardware, se integra una arquitectura basada en convoluciones separables (depthwise and pointwise convolution). A continuación, se define el bloque de convolución y la red estudiante (StudentNet) resultante, manteniendo el total de parámetros por debajo del límite de 60,000:

import torch.nn as nn

class SeparableConv2d(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0):
        super().__init__()
        self.depthwise = nn.Conv2d(
            in_channels, in_channels, kernel_size, 
            stride=stride, padding=padding, groups=in_channels
        )
        self.pointwise = nn.Conv2d(in_channels, out_channels, kernel_size=1)

    def forward(self, x):
        return self.pointwise(self.depthwise(x))

class CompactStudentNetwork(nn.Module):
    def __init__(self, num_classes=11):
        super().__init__()
        self.feature_extractor = nn.Sequential(
            nn.Conv2d(3, 4, 3), 
            nn.BatchNorm2d(4),
            nn.ReLU(inplace=True),    

            nn.Conv2d(4, 16, 3), 
            nn.BatchNorm2d(16),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2, 0),
            
            SeparableConv2d(16, 64, 3), 
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2, 0),

            SeparableConv2d(64, 64, 3), 
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, 3), 
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2, 0),
            
            SeparableConv2d(64, 64, 3), 
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            
            SeparableConv2d(64, 84, 3), 
            nn.BatchNorm2d(84),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2, 0),
            
            nn.AdaptiveAvgPool2d((1, 1))
        )
        
        self.classifier = nn.Linear(84, num_classes)

    def forward(self, x):
        features = self.feature_extractor(x)
        flattened = features.view(features.size(0), -1)
        return self.classifier(flattened)

En el contexto de la destilación de conocimiento (Knowledge Distillation), es fundamental implementar correctamente la función de pérdida que combina la Divergencia de Kullback-Leibler (KL) y la Entropía Cruzada. Un error común al calcular la divergencia KL es aplicar el softmax en la dimensión incorrecta o no transformar las probabilidades del estudiante a logaritmos, lo que resulta en un estancamiento del aprendizaje (por ejemplo, una precisión de solo 0.33 tras 50 épocas).

La implementación corregida y matemáticamente sólida de la función de pérdida es la siguiente:

import torch
import torch.nn.functional as F

def knowledge_distillation_loss(student_logits, teacher_logits, labels, temperature=2.0, alpha=0.5):
    # Calcular probabilidades suavizadas con temperatura
    student_log_probs = F.log_softmax(student_logits / temperature, dim=1)
    teacher_probs = F.softmax(teacher_logits / temperature, dim=1)
    
    # Pérdida por Divergencia de Kullback-Leibler
    kl_loss = F.kl_div(student_log_probs, teacher_probs, reduction="batchmean") * (temperature ** 2)
    
    # Pérdida estándar de Entropía Cruzada con las etiquetas reales
    ce_loss = F.cross_entropy(student_logits, labels)
    
    # Combinación ponderada de ambas pérdidas
    return alpha * kl_loss + (1 - alpha) * ce_loss

Adicionalmente, reestructurar las capas de la red estudiante permite optimizar el flujo de características y mejorar la convergencia. Se ajusta la arquitectura modificando los canales intermedios y la disposición de los bloques separables:

class RefinedStudentNetwork(nn.Module):
    def __init__(self, num_classes=11):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Conv2d(3, 4, 3), 
            nn.BatchNorm2d(4),
            nn.ReLU(inplace=True),    

            nn.Conv2d(4, 16, 3), 
            nn.BatchNorm2d(16),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2, 0),
            
            SeparableConv2d(16, 64, 3), 
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),

            SeparableConv2d(64, 32, 3), 
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2, 0),
            
            nn.Conv2d(32, 32, 3), 
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2, 0),
              
            SeparableConv2d(32, 128, 3), 
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),

            SeparableConv2d(128, 128, 3), 
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2, 0),
            
            SeparableConv2d(128, 84, 3), 
            nn.BatchNorm2d(84),
            nn.ReLU(inplace=True),
            
            nn.AdaptiveAvgPool2d((1, 1))
        )
        
        self.head = nn.Linear(84, num_classes)

    def forward(self, x):
        x = self.layers(x)
        x = torch.flatten(x, 1)
        return self.head(x)

Tras aplicar estas correcciones en la función de pérdida y ajustar la topología de la red, el modelo muestra una convergencia mucho más rápida, alcanzando una precisión del 0.35 desde la primera época de entrenamiento. Al completar el ciclo de 50 épocas, la red comprimida logra una puntuación final de 0.692, superando satisfactoriamente la línea base intermedia (medium baseline).