graphformer/models/finetune_model.py

"""
Fine-tuning Model for GraPHFormer

Supports three modes: image_only, tree_only, multimodal
Loads pretrained weights from CLIP-style training checkpoints
"""

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from .fusion import (
    CrossAttentionFusion, BiDirectionalCrossAttention,
    GatedFusion, CMF, MultiHeadCrossModalAttention
)


class ArcMarginProduct(nn.Module):
    """ArcFace: Additive Angular Margin Loss"""
    def __init__(self, in_features, out_features, s=30.0, m=0.50, easy_margin=False):
        super(ArcMarginProduct, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.s = s
        self.m = m
        self.easy_margin = easy_margin

        self.weight = nn.Parameter(torch.FloatTensor(out_features, in_features))
        nn.init.xavier_uniform_(self.weight)

        self.cos_m = np.cos(m)
        self.sin_m = np.sin(m)
        self.th = np.cos(np.pi - m)
        self.mm = np.sin(np.pi - m) * m

    def forward(self, input, label):
        cosine = F.linear(F.normalize(input), F.normalize(self.weight))
        sine = torch.sqrt(1.0 - torch.pow(cosine, 2))

        phi = cosine * self.cos_m - sine * self.sin_m

        if self.easy_margin:
            phi = torch.where(cosine > 0, phi, cosine)
        else:
            phi = torch.where(cosine > self.th, phi, cosine - self.mm)

        one_hot = torch.zeros(cosine.size(), device=input.device)
        one_hot.scatter_(1, label.view(-1, 1).long(), 1)

        output = (one_hot * phi) + ((1.0 - one_hot) * cosine)
        output *= self.s

        return output


class FineTuneModel(nn.Module):
    """Fine-tuning model with classification head"""
    def __init__(
        self,
        pretrained_model,
        num_classes,
        mode='multimodal',
        freeze_encoders=False,
        fusion_mode='concat',
        dropout=0.5,
        label_smoothing=0.0,
        use_projection=False,
        use_arcface=False,
        arcface_s=30.0,
        arcface_m=0.50,
        freeze_image_only=False
    ):
        """
        Args:
            pretrained_model: Pretrained CLIPModel
            num_classes: Number of classes for classification
            mode: 'image_only', 'tree_only', or 'multimodal'
            freeze_encoders: If True, freeze encoder weights
            fusion_mode: For multimodal - 'concat', 'add', 'cross_attention', 'bi_attention', 'gated', 'cmf', 'mhcma'
            dropout: Dropout rate for first layer
            label_smoothing: Label smoothing factor
            use_projection: If True, use projection heads from pretrained model
            use_arcface: If True, use ArcFace loss instead of CrossEntropy
            arcface_s: ArcFace scale parameter
            arcface_m: ArcFace margin parameter
            freeze_image_only: If True, freeze only image encoder
        """
        super(FineTuneModel, self).__init__()

        self.mode = mode
        self.fusion_mode = fusion_mode
        self.use_projection = use_projection
        self.use_arcface = use_arcface
        self.tree_encoder_type = pretrained_model.tree_encoder_type

        # Copy encoders from pretrained model
        if mode in ['tree_only', 'multimodal']:
            self.tree_encoder = pretrained_model.tree_encoder
            if use_projection:
                self.tree_projection = pretrained_model.tree_projection

        if mode in ['image_only', 'multimodal']:
            self.image_encoder = pretrained_model.image_encoder
            if use_projection:
                self.image_projection = pretrained_model.image_projection

        # Freeze encoders if requested
        if freeze_encoders:
            if mode in ['tree_only', 'multimodal']:
                for param in self.tree_encoder.parameters():
                    param.requires_grad = False
                if use_projection:
                    for param in self.tree_projection.parameters():
                        param.requires_grad = False

            if mode in ['image_only', 'multimodal']:
                for param in self.image_encoder.parameters():
                    param.requires_grad = False
                if use_projection:
                    for param in self.image_projection.parameters():
                        param.requires_grad = False

        # Freeze only image encoder
        if freeze_image_only and mode == 'multimodal':
            for param in self.image_encoder.parameters():
                param.requires_grad = False
            if use_projection:
                for param in self.image_projection.parameters():
                    param.requires_grad = False

        # Get embedding dimensions
        if use_projection:
            tree_embed_dim = pretrained_model.tree_projection[-1].out_features
            image_embed_dim = pretrained_model.image_projection[-1].out_features
        else:
            if mode in ['tree_only', 'multimodal']:
                tree_embed_dim = pretrained_model.tree_encoder.h_size
            else:
                tree_embed_dim = 0

            if mode in ['image_only', 'multimodal']:
                if hasattr(pretrained_model.image_encoder, 'encoder'):
                    if hasattr(pretrained_model.image_encoder.encoder, 'feat_dim'):
                        image_embed_dim = pretrained_model.image_encoder.encoder.feat_dim
                    else:
                        image_embed_dim = pretrained_model.image_encoder.encoder[-1].in_features
                elif hasattr(pretrained_model.image_encoder, 'feat_dim'):
                    image_embed_dim = pretrained_model.image_encoder.feat_dim
                else:
                    image_embed_dim = pretrained_model.tree_encoder.h_size
            else:
                image_embed_dim = 0

        # Setup fusion for multimodal
        if mode == 'multimodal':
            if fusion_mode == 'concat':
                fusion_dim = tree_embed_dim + image_embed_dim
            elif fusion_mode == 'add':
                fusion_dim = min(tree_embed_dim, image_embed_dim)
                if tree_embed_dim != image_embed_dim:
                    self.tree_dim_match = nn.Linear(tree_embed_dim, fusion_dim) if tree_embed_dim != fusion_dim else nn.Identity()
                    self.image_dim_match = nn.Linear(image_embed_dim, fusion_dim) if image_embed_dim != fusion_dim else nn.Identity()
            elif fusion_mode == 'cross_attention':
                fusion_dim = min(tree_embed_dim, image_embed_dim)
                if tree_embed_dim != image_embed_dim:
                    self.tree_dim_match = nn.Linear(tree_embed_dim, fusion_dim) if tree_embed_dim != fusion_dim else nn.Identity()
                    self.image_dim_match = nn.Linear(image_embed_dim, fusion_dim) if image_embed_dim != fusion_dim else nn.Identity()
                self.fusion_layer = CrossAttentionFusion(fusion_dim, num_heads=4)
            elif fusion_mode == 'bi_attention':
                fusion_dim = min(tree_embed_dim, image_embed_dim)
                if tree_embed_dim != image_embed_dim:
                    self.tree_dim_match = nn.Linear(tree_embed_dim, fusion_dim) if tree_embed_dim != fusion_dim else nn.Identity()
                    self.image_dim_match = nn.Linear(image_embed_dim, fusion_dim) if image_embed_dim != fusion_dim else nn.Identity()
                self.fusion_layer = BiDirectionalCrossAttention(fusion_dim, num_heads=4)
            elif fusion_mode == 'gated':
                fusion_dim = min(tree_embed_dim, image_embed_dim)
                if tree_embed_dim != image_embed_dim:
                    self.tree_dim_match = nn.Linear(tree_embed_dim, fusion_dim) if tree_embed_dim != fusion_dim else nn.Identity()
                    self.image_dim_match = nn.Linear(image_embed_dim, fusion_dim) if image_embed_dim != fusion_dim else nn.Identity()
                self.fusion_layer = GatedFusion(fusion_dim)
            elif fusion_mode == 'cmf':
                fusion_dim = min(tree_embed_dim, image_embed_dim)
                if tree_embed_dim != image_embed_dim:
                    self.tree_dim_match = nn.Linear(tree_embed_dim, fusion_dim) if tree_embed_dim != fusion_dim else nn.Identity()
                    self.image_dim_match = nn.Linear(image_embed_dim, fusion_dim) if image_embed_dim != fusion_dim else nn.Identity()
                self.fusion_layer = CMF(fusion_dim)
            elif fusion_mode == 'mhcma':
                fusion_dim = min(tree_embed_dim, image_embed_dim)
                if tree_embed_dim != image_embed_dim:
                    self.tree_dim_match = nn.Linear(tree_embed_dim, fusion_dim) if tree_embed_dim != fusion_dim else nn.Identity()
                    self.image_dim_match = nn.Linear(image_embed_dim, fusion_dim) if image_embed_dim != fusion_dim else nn.Identity()
                self.fusion_layer = MultiHeadCrossModalAttention(fusion_dim, num_heads=8)
            else:
                raise ValueError(f"Unknown fusion_mode: {fusion_mode}")
            input_dim = fusion_dim
        else:
            if mode == 'tree_only':
                input_dim = tree_embed_dim
            else:
                input_dim = image_embed_dim

        # Classification head
        if use_arcface:
            self.feature_extractor = nn.Sequential(
                nn.Dropout(dropout),
                nn.Linear(input_dim, input_dim // 2),
                nn.BatchNorm1d(input_dim // 2),
                nn.ReLU(),
                nn.Dropout(dropout * 0.7),
            )
            self.arcface = ArcMarginProduct(input_dim // 2, num_classes, s=arcface_s, m=arcface_m)
            self.classifier = None
        else:
            self.classifier = nn.Sequential(
                nn.Dropout(dropout),
                nn.Linear(input_dim, input_dim // 2),
                nn.BatchNorm1d(input_dim // 2),
                nn.ReLU(),
                nn.Dropout(dropout * 0.7),
                nn.Linear(input_dim // 2, num_classes)
            )
            self.feature_extractor = None
            self.arcface = None

        self.criterion = nn.CrossEntropyLoss(label_smoothing=label_smoothing)

    def encode_tree(self, batch):
        """Encode tree data"""
        tree_feats = self.tree_encoder(batch)

        if self.use_projection:
            tree_embed = self.tree_projection(tree_feats)
            return tree_embed
        else:
            return tree_feats

    def encode_image(self, images):
        """Encode image data"""
        image_feats = self.image_encoder(images)

        if self.use_projection:
            image_embed = self.image_projection(image_feats)
            return image_embed
        else:
            return image_feats

    def forward(self, batch, return_features=False):
        """
        Forward pass

        Args:
            batch: contains batch.graph, batch.feats, batch.images, batch.label
            return_features: if True, return embeddings along with logits
        Returns:
            loss: classification loss
            logits: (B, num_classes)
            features: (optional) embeddings
        """
        images = batch.images.cuda() if not batch.images.is_cuda else batch.images
        labels = batch.label.cuda() if not batch.label.is_cuda else batch.label

        if self.mode == 'tree_only':
            tree_embed = self.encode_tree(batch)
            tree_embed = F.normalize(tree_embed, dim=-1)
            features = tree_embed

        elif self.mode == 'image_only':
            image_embed = self.encode_image(images)
            image_embed = F.normalize(image_embed, dim=-1)
            features = image_embed

        else:  # multimodal
            tree_embed = self.encode_tree(batch)
            image_embed = self.encode_image(images)

            tree_embed = F.normalize(tree_embed, dim=-1)
            image_embed = F.normalize(image_embed, dim=-1)

            if self.fusion_mode == 'concat':
                features = torch.cat([tree_embed, image_embed], dim=1)
            elif self.fusion_mode == 'add':
                if hasattr(self, 'tree_dim_match'):
                    tree_embed = self.tree_dim_match(tree_embed)
                    image_embed = self.image_dim_match(image_embed)
                features = tree_embed + image_embed
            elif self.fusion_mode in ['gated', 'cmf', 'cross_attention', 'bi_attention', 'mhcma']:
                if hasattr(self, 'tree_dim_match'):
                    tree_embed = self.tree_dim_match(tree_embed)
                    image_embed = self.image_dim_match(image_embed)
                features = self.fusion_layer(tree_embed, image_embed)
            else:
                raise ValueError(f"Unknown fusion_mode: {self.fusion_mode}")

        # Classification
        if self.use_arcface:
            extracted_features = self.feature_extractor(features)
            logits = self.arcface(extracted_features, labels)
            loss = self.criterion(logits, labels)
        else:
            logits = self.classifier(features)
            loss = self.criterion(logits, labels)

        if return_features:
            return loss, logits, features
        return loss, logits

    def unfreeze_encoders(self):
        """Unfreeze encoder weights for full fine-tuning"""
        if self.mode in ['tree_only', 'multimodal'] and hasattr(self, 'tree_encoder'):
            for param in self.tree_encoder.parameters():
                param.requires_grad = True
            if self.use_projection and hasattr(self, 'tree_projection'):
                for param in self.tree_projection.parameters():
                    param.requires_grad = True

        if self.mode in ['image_only', 'multimodal'] and hasattr(self, 'image_encoder'):
            for param in self.image_encoder.parameters():
                param.requires_grad = True
            if self.use_projection and hasattr(self, 'image_projection'):
                for param in self.image_projection.parameters():
                    param.requires_grad = True