Tarea 7: Clasificador de imágenes de ocho formas con GoogLeNet

class GoogLeNet(nn.Module):
    def __init__(self, num_classes: int = 8, aux_logits: bool = True, ...):
        super().__init__()
        # Entrada adaptada para imagenes de 1 canal (escala de grises)
        self.conv1 = conv_block(1, 64, kernel_size=7, stride=2, padding=3)
        self.maxpool1 = nn.MaxPool2d(3, stride=2, ceil_mode=True)
        self.conv2 = conv_block(64, 64, kernel_size=1)
        self.conv3 = conv_block(64, 192, kernel_size=3, padding=1)
        self.maxpool2 = nn.MaxPool2d(3, stride=2, ceil_mode=True)

        # Se conserva unicamente el primer modulo Inception (inception3a)
        self.inception3a = inception_block(192, 64, 96, 128, 16, 32, 32)
        # Los modulos inception3b hasta inception5b fueron recortados

        if aux_logits:
            self.aux1 = inception_aux_block(256, num_classes, dropout=dropout_aux)
            self.aux2 = inception_aux_block(256, num_classes, dropout=dropout_aux)

        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.dropout = nn.Dropout(p=dropout)
        self.fc = nn.Linear(256, num_classes) # Clasificador lineal a 8 clases

class Inception(nn.Module):
    def __init__(self, in_channels: int, ch1x1: int, ch3x3red: int, ch3x3: int,
                 ch5x5red: int, ch5x5: int, pool_proj: int):
        super().__init__()
        self.branch1 = conv_block(in_channels, ch1x1, kernel_size=1)
        self.branch2 = nn.Sequential(
            conv_block(in_channels, ch3x3red, kernel_size=1),
            conv_block(ch3x3red, ch3x3, kernel_size=3, padding=1)
        )
        self.branch3 = nn.Sequential(
            conv_block(in_channels, ch5x5red, kernel_size=1),
            conv_block(ch5x5red, ch5x5, kernel_size=3, padding=1)
        )
        self.branch4 = nn.Sequential(
            nn.MaxPool2d(kernel_size=3, stride=1, padding=1, ceil_mode=True),
            conv_block(in_channels, pool_proj, kernel_size=1)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return torch.cat([self.branch1(x), self.branch2(x), self.branch3(x), self.branch4(x)], 1)

# Descarga de pesos preentrenados oficiales de PyTorch
pretrained_model = original_googlenet(weights='IMAGENET1K_V1')
pretrained_dict = pretrained_model.state_dict()
model_dict = nuevo_modelo.state_dict()

# Adaptar conv1 de 3 canales a 1 canal sumando sus filtros
if 'conv1.conv.weight' in pretrained_dict:
    peso_original = pretrained_dict['conv1.conv.weight'] # [64, 3, 7, 7]
    peso_modificado = peso_original.sum(dim=1, keepdim=True) # [64, 1, 7, 7]
    pretrained_dict['conv1.conv.weight'] = peso_modificado

# Cargar unicamente las capas presentes en nuestro modelo recortado
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict and v.shape == model_dict[k].shape}
model_dict.update(pretrained_dict)
nuevo_modelo.load_state_dict(model_dict)

# Pipeline de Transformaciones con torchvision v2
transforma = v2.Compose([
    v2.ToImage(),
    v2.RandomRotation(degrees=(0, 360)),
    v2.CenterCrop(size=58),
    v2.RandomHorizontalFlip(p=0.5),
    v2.RandomVerticalFlip(p=0.5),
    v2.ToDtype(torch.float32, scale=False)
])

# Custom Dataset para carga y normalizacion en escala de grises
class CustomImageDataset(Dataset):
    def __init__(self, annotations_file, img_dir, transform=None):
        self.img_labels = pd.read_csv(annotations_file, header=None)
        self.img_dir = img_dir
        self.transform = transform

    def __len__(self):
        return len(self.img_labels)

    def __getitem__(self, idx):
        img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0])
        image = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
        image = cv2.normalize(image, None, 0, 1.0, cv2.NORM_MINMAX, dtype=cv2.CV_32F)
        label = int(self.img_labels.iloc[idx, 1])
        if self.transform:
            image = self.transform(image)
        return image, label

# Congelar capas base, descongelar inception3a y clasificador final
for name, param in nuevo_modelo.named_parameters():
    if "fc" not in name and "aux" not in name and "inception3a" not in name:
        param.requires_grad = False
    else:
        param.requires_grad = True

# Optimizador SGD con Nesterov Momentum
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(
    filter(lambda p: p.requires_grad, nuevo_modelo.parameters()),
    lr=0.1,
    momentum=0.9,
    nesterov=True
)