Novel and Emerging ML Techniques Questions

import torch
import copy
from collections import OrderedDict

def maml_inner_update(model, x_s, y_s, loss_fn, inner_lr):
    """
    Perform one MAML inner-loop update and return an OrderedDict of adapted params.
    Gradients are computed with create_graph=True so outer-loop gradients can flow.
    """
    # forward
    preds = model(x_s)
    loss = loss_fn(preds, y_s)

    # get parameters as (name, param) list to preserve order
    params = OrderedDict(model.named_parameters())
    # compute gradients w.r.t. parameters (create_graph=True for higher-order grads)
    grads = torch.autograd.grad(loss, params.values(), create_graph=True, allow_unused=True)

    adapted = OrderedDict()
    for (name, p), g in zip(params.items(), grads):
        if g is None:
            # if param did not affect loss, keep it unchanged
            adapted[name] = p
        else:
            adapted[name] = p - inner_lr * g

    return adapted

# Example to load adapted params into a cloned model:
# cloned = copy.deepcopy(model)
# cloned.load_state_dict(adapted, strict=False)  # works if adapted covers same keys
# now use cloned for validation; gradients will flow back through adapted params.

import torch
import torch.nn as nn
import torchvision.models as models
from typing import Optional

def prepare_finetune_model(model_name: str, num_classes: int, freeze_until: Optional[str]) -> torch.nn.Module:
    """
    Load a pretrained model from torchvision, freeze parameters up to (and including)
    the module named `freeze_until` (if provided), handle BatchNorm correctly, replace
    the classifier head to match num_classes, and return the model ready for training.

    freeze_until: None -> don't freeze any layers
                  str  -> module name (as in model.named_modules()) up to which to freeze (inclusive)
    """
    # load pretrained model
    if not hasattr(models, model_name):
        raise ValueError(f"Unknown model_name: {model_name}")
    # Note: depending on torchvision version, pretrained arg may differ (weights=...). This uses classic API.
    model = getattr(models, model_name)(pretrained=True)

    # Freeze logic: iterate named_modules top-down, freeze until we hit freeze_until (inclusive).
    # If freeze_until is None, we don't freeze any layers.
    freeze = False if freeze_until is None else True
    stopped = freeze_until is None  # if None, we've "already stopped" freezing
    for mod_name, module in model.named_modules():
        if stopped:
            break if False else None  # no-op to keep clarity; we use freeze flag below
        # We'll manage freezing by toggling when we hit the target name.
        # However we need to apply freezing to parameters - we do that by walking parameters of each module.
        # (handled below)
        if freeze:
            # apply freezing to parameters of this module
            for p in module.parameters(recurse=False):
                p.requires_grad = False
            # If module is BatchNorm, set to eval to use running stats and ensure its affine params don't update
            if isinstance(module, (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d)):
                module.eval()
                for name, p in module.named_parameters(recurse=False):
                    p.requires_grad = False
        # Check if current module is the stop marker
        if mod_name == freeze_until:
            stopped = True
            freeze = False  # stop freezing subsequent modules

    # Edge case: if freeze_until was provided but not found, raise to avoid silent mistakes
    if freeze_until is not None and not stopped:
        raise ValueError(f"freeze_until='{freeze_until}' not found among model.named_modules()")

    # Replace classifier head for common model families
    # ResNet-like
    if hasattr(model, 'fc') and isinstance(model.fc, nn.Module):
        in_feats = model.fc.in_features
        model.fc = nn.Linear(in_feats, num_classes)
    # Models with classifier as Sequential or Linear (MobileNet, EfficientNet, VGG, AlexNet)
    elif hasattr(model, 'classifier') and isinstance(model.classifier, (nn.Sequential, nn.Module)):
        # Many classifier implementations are either a single Linear or a Sequential ending in Linear
        if isinstance(model.classifier, nn.Sequential):
            # replace last linear layer
            # find last linear layer index
            last_lin_idx = None
            for i, layer in reversed(list(enumerate(model.classifier))):
                if isinstance(layer, nn.Linear):
                    last_lin_idx = i
                    break
            if last_lin_idx is None:
                # fallback: replace entire classifier with a single linear
                # try to infer in_features from first linear found in children
                in_feats = None
                for layer in model.classifier:
                    if isinstance(layer, nn.Linear):
                        in_feats = layer.in_features
                        break
                if in_feats is None:
                    raise RuntimeError("Cannot infer classifier in_features to replace head")
                model.classifier = nn.Linear(in_feats, num_classes)
            else:
                in_feats = model.classifier[last_lin_idx].in_features
                model.classifier[last_lin_idx] = nn.Linear(in_feats, num_classes)
        else:
            # single module, try to replace if linear
            if isinstance(model.classifier, nn.Linear):
                in_feats = model.classifier.in_features
                model.classifier = nn.Linear(in_feats, num_classes)
            else:
                # replace with a simple linear layer if unsure
                raise RuntimeError("Unhandled classifier shape; please adapt replacement logic")
    else:
        raise RuntimeError("Unhandled model head type; add logic for this architecture")

    return model