third_party/InternVideo/InternVideo2/multi_modality/tasks/shared_utils.py

import copy
import logging
import os
import os.path as osp
import io

try:
    import deepspeed
except Exception as e:
    print(e)
    print("deepspeed is not installed!!!")
    
from os.path import join


try:
    from petrel_client.client import Client
except:
    Client = None

import torch
from torch.utils.data import ConcatDataset, DataLoader

from dataset.resample_concat_dataset import ResampleConcatDataset
from models.backbones.internvideo2.pos_embed import interpolate_pos_embed_internvideo2_new
from models.backbones.bert.tokenization_bert import BertTokenizer
from utils.optimizer import create_optimizer
from utils.scheduler import create_scheduler
from utils.distributed import get_rank

logger = logging.getLogger(__name__)


def get_media_types(datasources):
    """get the media types for for all the dataloaders.

    Args:
        datasources (List): List of dataloaders or datasets.

    Returns: List. The media_types.

    """
    if isinstance(datasources[0], DataLoader):
        datasets = [dataloader.dataset for dataloader in datasources]
    else:
        datasets = datasources
    media_types = [
        dataset.datasets[0].media_type
        if isinstance(dataset, ConcatDataset) or isinstance(dataset, ResampleConcatDataset)
        else dataset.media_type
        for dataset in datasets
    ]

    return media_types


def setup_model(
    config, model_cls, add_decoder=False, pretrain=False, find_unused_parameters=False
):
    logger.info("Creating model")
    config = copy.deepcopy(config)

    if "bert" in config.model.text_encoder.name:
        logger.info(f"Using BertTokenizer: {config.model.text_encoder.pretrained}!")
        tokenizer = BertTokenizer.from_pretrained(config.model.text_encoder.pretrained, local_files_only=True)
        model = model_cls(config=config, tokenizer=tokenizer, is_pretrain=pretrain)
    else:
        model = model_cls(config=config, is_pretrain=pretrain)
        tokenizer = model.tokenizer
        logger.info(f"Using model.tokenizer: {tokenizer}!")

    
    if config.get('compile_model', False):
        torch.set_float32_matmul_precision('high')
        model = torch.compile(model)

    model = model.to(torch.device(config.device))
    model_without_ddp = model


    if hasattr(config, "deepspeed") and config.deepspeed.enable:
        # We move this to the back
        optimizer_params = create_optimizer(config.optimizer, model, return_group=True)
        scheduler = None
        scaler = None
    else:
        if config.distributed:
            model = torch.nn.parallel.DistributedDataParallel(
                model,
                device_ids=[config.gpu],
                find_unused_parameters=find_unused_parameters,  # `False` for image-only task
            )

        optimizer = create_optimizer(config.optimizer, model)
        scaler = torch.cuda.amp.GradScaler(enabled=config.use_half_precision)  # This is never used actually if we fixed bf16
        scheduler = create_scheduler(config.scheduler, optimizer)


    start_epoch = 0
    global_step = 0

    # auto resume the latest checkpoint
    if config.get("auto_resume", False):
        logger.info("Auto resuming")
        model_latest = join(config.output_dir, "ckpt_latest.pth")
        model_best = join(config.output_dir, "ckpt_best.pth")
        large_num = -1
        for p in os.listdir(config.output_dir):
            if 'ckpt' in p:
                num = p.split('_')[1].split('.')[0]
                if str.isnumeric(num):
                    if int(num) > large_num:
                        large_num = int(num)
        if large_num != -1:
            model_latest = join(config.output_dir, f"ckpt_{large_num:02d}.pth")
        if osp.isfile(model_latest):
            config.pretrained_path = model_latest
            config.resume = True
        elif osp.isfile(model_best):
            config.pretrained_path = model_best
            config.resume = True
        else:
            logger.info(f"Not found checkpoint in {config.output_dir}")


    if (config.pretrained_path.strip() and (osp.isfile(config.pretrained_path)) or "s3://" in config.pretrained_path):
        if Client is not None:
            client = Client()
            with io.BytesIO(client.get(config.pretrained_path)) as buffer:
                checkpoint = torch.load(buffer, map_location="cpu")
        else:
            checkpoint = torch.load(config.pretrained_path, map_location="cpu")
        logger.info(f"Loading checkpoint from {config.pretrained_path}")
        try:
            if "model" in checkpoint.keys():
                state_dict = checkpoint["model"]
            else:
                state_dict = checkpoint["module"] # This is a deepspeed stage 1 model
        except:  
            state_dict = checkpoint

        if config.get('origin_num_frames', None) is not None:
            logger.info(f"interpolate_pos_embed_internvideo2 (origin_num_frames={config.origin_num_frames})!!!")
            a = len(state_dict)
            interpolate_pos_embed_internvideo2_new(state_dict, model_without_ddp.vision_encoder, orig_t_size=config.origin_num_frames)
            assert a == len(state_dict), state_dict.keys()

        if config.resume:
            assert not (hasattr(config, "deepspeed") and config.deepspeed.enable), "Deepspeed should run here!!!"
            optimizer.load_state_dict(checkpoint["optimizer"])
            scheduler.load_state_dict(checkpoint["scheduler"])
            scaler.load_state_dict(checkpoint["scaler"])
            if 'local_step' in checkpoint.keys():
                start_epoch = checkpoint['epoch']
            else:
                start_epoch = checkpoint["epoch"] + 1
            global_step = checkpoint["global_step"]

        elif not pretrain:  # downstream init from pretrained ckpt 


            if not config.evaluate or config.get("zero_shot", False):  # finetuning from a pretrained weights.
                if add_decoder:
                    logger.info("Init new decoder with encoder!!!")
                for key in list(state_dict.keys()):
                    if "text_encoder.bert" in key:
                        encoder_key = key.replace("bert.", "") 
                        state_dict[encoder_key] = state_dict[key]
                        if not add_decoder:
                            del state_dict[key]

                    # init text decoder as multimodal encoder (last 6 layers of model.text_encoder)
                    # only for generation tasks like VQA
                    if add_decoder and "text_encoder.bert" in key:
                        if "layer" in key:
                            encoder_keys = key.split(".")
                            layer_num = int(encoder_keys[4])
                            if layer_num < config.model.text_encoder.fusion_layer:
                                del state_dict[key]
                                continue
                            else:
                                decoder_layer_num = layer_num - config.model.text_encoder.fusion_layer
                                encoder_keys[4] = str(decoder_layer_num)
                                encoder_key = ".".join(encoder_keys)
                        else:
                            encoder_key = key
                        decoder_key = encoder_key.replace("text_encoder", "text_decoder")
                        state_dict[decoder_key] = state_dict[key]
                        del state_dict[key]


        msg = model_without_ddp.load_state_dict(state_dict, strict=False)
        logger.info(msg)
        logger.info(f"Loaded checkpoint from {config.pretrained_path}")
    else:
        if not config.resume:
            assert not config.evaluate, "No available pretrained checkpoint provided!!!"
            assert config.pretrained_path == "", config.pretrained_path
            logger.warning("No available pretrained checkpoint provided, training from scratch.")


    if hasattr(config, "deepspeed") and config.deepspeed.enable:
        logger.info(f'Use deepspeed to initialize model (resume={config.resume}) !!!')
        model = model_without_ddp

        model, optimizer, _, _ = deepspeed.initialize(
            args=config, model=model, model_parameters=optimizer_params, dist_init_required=not config.distributed,
            lr_scheduler=lambda opt: create_scheduler(config.scheduler, opt)
        )


        if config.resume:
            logger.info(f'Resume deepspeed ckpt from {config.output_dir}, tag={config.pretrained_path}, load_module_strict={config.get("load_module_strict", True)}, load_lr_scheduler_states={config.get("load_lr_scheduler_states", True)}!!!')
            _, client_states = model.load_checkpoint(config.output_dir, tag=config.pretrained_path,  load_module_strict=config.get("load_module_strict", True), load_lr_scheduler_states=config.get("load_lr_scheduler_states", True))
            logger.info(client_states)
            if 'local_step' in client_states.keys():
                start_epoch = client_states['epoch']
            else:
                start_epoch = client_states['epoch'] + 1
            global_step = client_states['global_step']


    logger.info(f"Cuda memory after create model: {torch.cuda.memory_allocated() // 1024**2}M, Max mem: {torch.cuda.max_memory_allocated() // 1024**2}M start_epoch={start_epoch}, global_step={global_step}")
    print(f"\033[31m Cuda memory after create model: {torch.cuda.memory_allocated() // 1024**2}M, Max mem: {torch.cuda.max_memory_allocated() // 1024**2}M start_epoch={start_epoch}, global_step={global_step}\033[0m")

    return (
        model,
        model_without_ddp,
        optimizer,
        scheduler,
        scaler,
        tokenizer,
        start_epoch,
        global_step,
    )