loralib/utils.py

import os

import torch
import torch.nn as nn
from typing import Dict

from .layers import LoRALayer, AttentionLoRA, BertAttentionLoRA
from timm.models.vision_transformer import Attention
from transformers.models.bert.modeling_bert import BertAttention


INDEX_POSITIONS_TEXT = {
    'top1': [11],
    'top2': [10, 11],
    'top3': [9, 10, 11],
    'bottom': [0, 1, 2, 3],
    'mid': [4, 5, 6, 7],
    'up': [8, 9, 10, 11],
    'half-up': [6, 7, 8, 9, 10, 11],
    'half-bottom': [0, 1, 2, 3, 4, 5],
    'all': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]}


INDEX_POSITIONS_VISION = {
    'top': [11],
    'top3': [9, 10, 11],
    'bottom': [0, 1, 2, 3],
    'mid': [4, 5, 6, 7],
    'up': [8, 9, 10, 11],
    'half-up': [6, 7, 8, 9, 10, 11],
    'half-bottom': [0, 1, 2, 3, 4, 5],
    'all': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11],
}


def mark_only_lora_as_trainable(model: nn.Module, bias: str = 'none') -> None:
    for n, p in model.named_parameters():
        if 'lora_' not in n:
            p.requires_grad = False
    if bias == 'none':
        return
    elif bias == 'all':
        for n, p in model.named_parameters():
            if 'bias' in n:
                p.requires_grad = True
    elif bias == 'lora_only':
        for m in model.modules():
            if isinstance(m, LoRALayer) and \
                    hasattr(m, 'bias') and \
                    m.bias is not None:
                m.bias.requires_grad = True
    else:
        raise NotImplementedError


def lora_state_dict(model: nn.Module, bias: str = 'none') -> Dict[str, torch.Tensor]:
    my_state_dict = model.state_dict()
    if bias == 'none':
        return {k: my_state_dict[k] for k in my_state_dict if 'lora_' in k}
    elif bias == 'all':
        return {k: my_state_dict[k] for k in my_state_dict if 'lora_' in k or 'bias' in k}
    elif bias == 'lora_only':
        to_return = {}
        for k in my_state_dict:
            if 'lora_' in k:
                to_return[k] = my_state_dict[k]
                bias_name = k.split('lora_')[0]+'bias'
                if bias_name in my_state_dict:
                    to_return[bias_name] = my_state_dict[bias_name]
        return to_return
    else:
        raise NotImplementedError


def get_lora_parameters(model, bias='none'):
    params = []
    for name, param in model.named_parameters():
        if bias == 'none':
            if 'lora_' in name:
                params.append(param)
        elif bias == 'all':
            if 'lora_' in name or 'bias' in name:
                params.append(param)
        elif bias == 'lora_only':
            if 'lora_' in name:
                params.append(param)
                bias_name = name.split('lora_')[0] + 'bias'
                if bias_name in model.state_dict():
                    bias_param = dict(model.named_parameters())[bias_name]
                    params.append(bias_param)
        else:
            raise NotImplementedError
    return params


def apply_lora(args, clip_model):
    list_lora_layers = []
    indices = INDEX_POSITIONS_TEXT[args.position]
    text_encoder = clip_model.text.transformer.encoder
    for i, block in enumerate(text_encoder.layer):
        if i in indices:
            for name, submodule in block.named_children():
                if isinstance(submodule, BertAttention):
                    new_multi_head_lora = BertAttentionLoRA(
                        submodule, enable_lora=args.params, r=args.r, lora_alpha=args.alpha, dropout_rate=args.dropout_rate, seed=args.seed)
                    setattr(block, name, new_multi_head_lora)
                    list_lora_layers.append(new_multi_head_lora)

    indices = INDEX_POSITIONS_VISION[args.position]
    vision_encoder = clip_model.visual.trunk
    for i, block in enumerate(vision_encoder.blocks):
        if i in indices:
            for name, submodule in block.named_children():
                if isinstance(submodule, Attention):
                    new_multi_head_lora = AttentionLoRA(
                        submodule, enable_lora=args.params, r=args.r, lora_alpha=args.alpha, dropout_rate=args.dropout_rate, seed=args.seed)
                    setattr(block, name, new_multi_head_lora)
                    list_lora_layers.append(new_multi_head_lora)
    return list_lora_layers


def save_lora(args, list_lora_layers, loss_fn, msg, save_dir):
    weights = {}
    for i, layer in enumerate(list_lora_layers):
        layer_weights = {}
        if 'q' in args.params:
            layer_weights['q_proj'] = {
                'w_lora_A': layer.q_proj.w_lora_A.data,
                'w_lora_B': layer.q_proj.w_lora_B.data
            }
        if 'k' in args.params:
            layer_weights['k_proj'] = {
                'w_lora_A': layer.k_proj.w_lora_A.data,
                'w_lora_B': layer.k_proj.w_lora_B.data
            }
        if 'v' in args.params:
            layer_weights['v_proj'] = {
                'w_lora_A': layer.v_proj.w_lora_A.data,
                'w_lora_B': layer.v_proj.w_lora_B.data
            }
        if 'o' in args.params:
            layer_weights['proj'] = {
                'w_lora_A': layer.proj.w_lora_A.data,
                'w_lora_B': layer.proj.w_lora_B.data
            }

        weights[f'layer_{i}'] = layer_weights
    
    if args.loss_type == 'clip_loss_ace_hgnn':
        weights['img_edge_adapter'] = loss_fn.img_edge_adapter.state_dict()
        weights['img_node_adapter'] = loss_fn.img_node_adapter.state_dict()
        weights['text_edge_adapter'] = loss_fn.text_edge_adapter.state_dict()
        weights['text_node_adapter'] = loss_fn.text_node_adapter.state_dict()
    
    if args.learnable_logit_scale:
        weights['logit_scale'] = loss_fn.logit_scale.data.cpu()

    metadata = {
        'r': args.r,
        'topk': args.topk,
        'params': args.params,
        'position': args.position,
        'loss_type' : args.loss_type,
    }

    save_data = {
        'weights': weights,
        'metadata': metadata
    }

    save_path = f'{save_dir}/{args.filename}_{msg}.pt'
    torch.save(save_data, save_path)
    print(f'LoRA weights saved to {save_path}')

def load_model(args, list_lora_layers, device, loss_fn=None):

    if not os.path.exists(args.load_path):
        raise FileNotFoundError(f'File {args.load_path} does not exist.')

    loaded_data = torch.load(args.load_path, map_location=device)

    weights = loaded_data['weights']
    for i, layer in enumerate(list_lora_layers):
        layer_weights = weights[f'layer_{i}']
        if 'q' in args.params and 'q_proj' in layer_weights:
            layer.q_proj.w_lora_A.data.copy_(
                layer_weights['q_proj']['w_lora_A'])
            layer.q_proj.w_lora_B.data.copy_(
                layer_weights['q_proj']['w_lora_B'])
        if 'k' in args.params and 'k_proj' in layer_weights:
            layer.k_proj.w_lora_A.data.copy_(
                layer_weights['k_proj']['w_lora_A'])
            layer.k_proj.w_lora_B.data.copy_(
                layer_weights['k_proj']['w_lora_B'])
        if 'v' in args.params and 'v_proj' in layer_weights:
            layer.v_proj.w_lora_A.data.copy_(
                layer_weights['v_proj']['w_lora_A'])
            layer.v_proj.w_lora_B.data.copy_(
                layer_weights['v_proj']['w_lora_B'])
        if 'o' in args.params and 'proj' in layer_weights:
            layer.proj.w_lora_A.data.copy_(layer_weights['proj']['w_lora_A'])
            layer.proj.w_lora_B.data.copy_(layer_weights['proj']['w_lora_B'])

    if args.loss_type == 'clip_loss_ace_hgnn':
        loss_fn.img_edge_adapter.load_state_dict(weights['img_edge_adapter'])
        loss_fn.img_node_adapter.load_state_dict(weights['img_node_adapter'])
        loss_fn.text_edge_adapter.load_state_dict(weights['text_edge_adapter'])
        loss_fn.text_node_adapter.load_state_dict(weights['text_node_adapter'])

    if args.learnable_logit_scale:
        loss_fn.logit_scale.data.copy_(weights['logit_scale'])

    print(f'LoRA weights loaded from {args.load_path}')