code/base.py

import os.path as osp
import math
import abc
from torch.utils.data import DataLoader
import torch.optim
import torchvision.transforms as transforms
from timer import Timer
from logger import colorlogger
from torch.nn.parallel.data_parallel import DataParallel
from config import cfg
from SMPLer_X import get_model
from dataset import MultipleDatasets
# ddp
import torch.distributed as dist
from torch.utils.data import DistributedSampler
import torch.utils.data.distributed
from utils.distribute_utils import (
    get_rank, is_main_process, time_synchronized, get_group_idx, get_process_groups
)
from mmcv.runner import get_dist_info
import loralib as lora 
from pdb import set_trace


# dynamic dataset import
for i in range(len(cfg.trainset_3d)):
    exec('from ' + cfg.trainset_3d[i] + ' import ' + cfg.trainset_3d[i])
for i in range(len(cfg.trainset_2d)):
    exec('from ' + cfg.trainset_2d[i] + ' import ' + cfg.trainset_2d[i])
for i in range(len(cfg.trainset_humandata)):
    exec('from ' + cfg.trainset_humandata[i] + ' import ' + cfg.trainset_humandata[i])
exec('from ' + cfg.testset + ' import ' + cfg.testset)


class Base(object):
    __metaclass__ = abc.ABCMeta

    def __init__(self, log_name='logs.txt'):
        self.cur_epoch = 0

        # timer
        self.tot_timer = Timer()
        self.gpu_timer = Timer()
        self.read_timer = Timer()

        # logger
        self.logger = colorlogger(cfg.log_dir, log_name=log_name)

    @abc.abstractmethod
    def _make_batch_generator(self):
        return

    @abc.abstractmethod
    def _make_model(self):
        return


class Trainer(Base):
    def __init__(self, distributed=False, gpu_idx=None, use_lora=False):
        super(Trainer, self).__init__(log_name='train_logs.txt')
        self.distributed = distributed
        self.gpu_idx = gpu_idx
        self.use_lora = use_lora

    def get_optimizer(self, model):
        normal_param = []
        special_param = []
        for module in model.module.special_trainable_modules:
            special_param += list(module.parameters())
            # print(module)
        for module in model.module.trainable_modules:
            normal_param += list(module.parameters())
        # self.logger.info(f"N-{self.gpu_idx}, {normal_param}")
        # self.logger.info("S", special_param)
        optim_params = [
            {  # add normal params first
                'params': normal_param,
                'lr': cfg.lr
            },
            {
                'params': special_param,
                'lr': cfg.lr * cfg.lr_mult
            },
        ]
        optimizer = torch.optim.Adam(optim_params, lr=cfg.lr)
        return optimizer

    def save_model(self, state, epoch):
        file_path = osp.join(cfg.model_dir, 'snapshot_{}.pth.tar'.format(str(epoch)))

        # do not save smplx layer weights
        dump_key = []
        for k in state['network'].keys():
            if 'smplx_layer' in k:
                dump_key.append(k)
        for k in dump_key:
            state['network'].pop(k, None)

        torch.save(state, file_path)
        self.logger.info("Write snapshot into {}".format(file_path))

    def load_model(self, model, optimizer):
        

        if cfg.pretrained_model_path is not None:
            ckpt_path = cfg.pretrained_model_path
            

            ckpt = torch.load(ckpt_path, map_location=torch.device('cpu')) # solve CUDA OOM error in DDP

            state = ckpt['network']

            # # NO-ROI
            # drop_prefixes = ('module.hand_position_net.', 'hand_position_net.',
            #                 'module.hand_regressor.', 'hand_regressor.',)
            # drop_keys = [k for k in list(state.keys()) if k.startswith(drop_prefixes)]
            # for k in drop_keys:
            #     del state[k]
            # self.logger.info(f"[load_model] drop hand_position_net keys: {len(drop_keys)}")


            model.load_state_dict(ckpt['network'], strict=False)
            self.logger.info('Load checkpoint from {}'.format(ckpt_path))
            if not hasattr(cfg, 'start_over') or cfg.start_over:
                start_epoch = 0
                self.logger.info("Start over: reinitializing cam_out parameters.")

                with torch.no_grad():
                    for name, param in model.named_parameters():
                        if 'cam_out' in name:
                            self.logger.info(f"Reinitializing {name}")
                            if 'weight' in name:
                                torch.nn.init.kaiming_normal_(param)
                            elif 'bias' in name:
                                torch.nn.init.constant_(param, 0.0)
            else:
                optimizer.load_state_dict(ckpt['optimizer'])
                start_epoch = ckpt['epoch'] + 1 
                self.logger.info(f'Load optimizer, start from{start_epoch}')
        else:
            start_epoch = 0

        if getattr(cfg, 'token_decoder_ckpt_path', None): 
            td_ckpt_path = cfg.token_decoder_ckpt_path
            td_ckpt = torch.load(td_ckpt_path, map_location=torch.device('cpu'))
            # set_trace()
            self.logger.info(f'Load token-decoder ckpt (smpl_head only) from {td_ckpt_path}')
            td_state = td_ckpt.get('state_dict', td_ckpt)
            if any(k.startswith('module.') for k in td_state.keys()):
                td_state = {k.replace('module.', '', 1): v for k, v in td_state.items()}
            token_decoder_prefix = 'token_decoder.'
            mapped = {}
            for k, v in td_state.items():
                if k.startswith('smpl_head.'):
                    subkey = k[len('smpl_head.'):]                 # ex) 'transformer.pos_embedding' ...
                    mapped[token_decoder_prefix + subkey] = v 
            needs_module_prefix = hasattr(model, 'module')
            if needs_module_prefix:
                mapped = {('module.' + k): v for k, v in mapped.items()}

            missing, unexpected = model.load_state_dict(mapped, strict=False)
            self.logger_info(
                f"token_decoder override done. Missing: {len(missing)}, Unexpected: {len(unexpected)}"
            )
            if missing:
                self.logger_info(f"  some missing (expected when 다른 파트 키 포함 안 함): {missing[:8]}{' ...' if len(missing)>8 else ''}")
            if unexpected:
                self.logger_info(f"  unexpected keys (확인 필요): {unexpected[:8]}{' ...' if len(unexpected)>8 else ''}")
            # set_trace()

                

        return start_epoch, model, optimizer

    def get_lr(self):
        for g in self.optimizer.param_groups:
            cur_lr = g['lr']
        return cur_lr

    # def _make_batch_generator(self):
    #     # data load and construct batch generator
    #     self.logger_info("Creating dataset...")
    #     trainset3d_loader = []
    #     for i in range(len(cfg.trainset_3d)):
    #         trainset3d_loader.append(eval(cfg.trainset_3d[i])(transforms.ToTensor(), "train"))
    #     trainset2d_loader = []
    #     for i in range(len(cfg.trainset_2d)):
    #         trainset2d_loader.append(eval(cfg.trainset_2d[i])(transforms.ToTensor(), "train"))
    #     trainset_humandata_loader = []
    #     for i in range(len(cfg.trainset_humandata)):
    #         trainset_humandata_loader.append(eval(cfg.trainset_humandata[i])(transforms.ToTensor(), "train"))
        
    #     data_strategy = getattr(cfg, 'data_strategy', None)
    #     if data_strategy == 'concat':
    #         print("Using [concat] strategy...")
    #         trainset_loader = MultipleDatasets(trainset3d_loader + trainset2d_loader + trainset_humandata_loader, 
    #                                             make_same_len=False, verbose=True)
    #     elif data_strategy == 'balance':
    #         total_len = getattr(cfg, 'total_data_len', 'auto')
    #         print(f"Using [balance] strategy with total_data_len : {total_len}...")
    #         trainset_loader = MultipleDatasets(trainset3d_loader + trainset2d_loader + trainset_humandata_loader, 
    #                                              make_same_len=True, total_len=total_len, verbose=True)
    #     else:
    #         # original strategy implementation
    #         valid_loader_num = 0
    #         if len(trainset3d_loader) > 0:
    #             trainset3d_loader = [MultipleDatasets(trainset3d_loader, make_same_len=False)]
    #             valid_loader_num += 1
    #         else:
    #             trainset3d_loader = []
    #         if len(trainset2d_loader) > 0:
    #             trainset2d_loader = [MultipleDatasets(trainset2d_loader, make_same_len=False)]
    #             valid_loader_num += 1
    #         else:
    #             trainset2d_loader = []
    #         if len(trainset_humandata_loader) > 0:
    #             trainset_humandata_loader = [MultipleDatasets(trainset_humandata_loader, make_same_len=False)]
    #             valid_loader_num += 1

    #         if valid_loader_num > 1:
    #             trainset_loader = MultipleDatasets(trainset3d_loader + trainset2d_loader + trainset_humandata_loader, make_same_len=True)
    #         else:
    #             trainset_loader = MultipleDatasets(trainset3d_loader + trainset2d_loader + trainset_humandata_loader, make_same_len=False)

    #     self.itr_per_epoch = math.ceil(len(trainset_loader) / cfg.num_gpus / cfg.train_batch_size)

    #     if self.distributed:
    #         self.logger_info(f"Total data length {len(trainset_loader)}.")
    #         rank, world_size = get_dist_info()
    #         self.logger_info("Using distributed data sampler.")
            
    #         sampler_train = DistributedSampler(trainset_loader, world_size, rank, shuffle=True)
    #         self.batch_generator = DataLoader(dataset=trainset_loader, batch_size=cfg.train_batch_size,
    #                                       shuffle=False, num_workers=cfg.num_thread, sampler=sampler_train,
    #                                       pin_memory=True, persistent_workers=True if cfg.num_thread > 0 else False, drop_last=True)
    #     else:
    #         self.batch_generator = DataLoader(dataset=trainset_loader, batch_size=cfg.num_gpus * cfg.train_batch_size,
    #                                       shuffle=True, num_workers=cfg.num_thread,
    #                                       pin_memory=True, drop_last=True)

    def _make_batch_generator(self):
        # data load and construct batch generator
        self.logger_info("Creating dataset...")
        trainset3d_loader = []
        for i in range(len(cfg.trainset_3d)):
            trainset3d_loader.append(eval(cfg.trainset_3d[i])(transforms.ToTensor(), "train"))
        trainset2d_loader = []
        for i in range(len(cfg.trainset_2d)):
            trainset2d_loader.append(eval(cfg.trainset_2d[i])(transforms.ToTensor(), "train"))
        trainset_humandata_loader = []
        for i in range(len(cfg.trainset_humandata)):
            trainset_humandata_loader.append(eval(cfg.trainset_humandata[i])(transforms.ToTensor(), "train"))

        # -------------------- Only if ID loss is enabled: build & apply shared id_roster --------------------
        use_id_loss = (float(getattr(cfg, "id_adv_lambda", 0.0)) > 0.0
                       and float(getattr(cfg, "id_loss_weight", 0.0)) > 0.0)
        if use_id_loss:
            candidates = [ds for ds in trainset_humandata_loader if hasattr(ds, "datalist")]
            all_ids = []
            for ds in candidates:
                # collect from datalist if available
                ids = [it.get("id_str") for it in getattr(ds, "datalist", []) if it.get("id_str") is not None]
                # fallback: some datasets expose IDs roster
                if not ids and hasattr(ds, "IDs"):
                    ids = list(getattr(ds, "IDs"))
                all_ids.extend(ids)

            shared_roster = sorted(set(all_ids))
            if len(shared_roster) == 0:
                print("[ID] WARNING: no id_str found in datasets. Disable ID loss by setting cfg.num_subjects=0.")
                from config import cfg as _cfg
                _cfg.num_subjects = 0
            else:
                id2idx = {sid: i for i, sid in enumerate(shared_roster)}
                for ds in candidates:
                    ds.IDs = list(shared_roster)
                    ds.ID_label2idx = dict(id2idx)
                    ds.num_ids = len(shared_roster)
                    for it in ds.datalist:
                        sid = it.get("id_str")
                        if sid is not None:
                            it["id_idx"] = id2idx[sid]
                from config import cfg as _cfg
                _cfg.num_subjects = len(shared_roster)
                print(f"[ID] Shared roster({_cfg.num_subjects}): {shared_roster}")
        # ----------------------------------------------------------------------------------------------------

        data_strategy = getattr(cfg, 'data_strategy', None)
        if data_strategy == 'concat':
            print("Using [concat] strategy...")
            trainset_loader = MultipleDatasets(trainset3d_loader + trainset2d_loader + trainset_humandata_loader, 
                                                make_same_len=False, verbose=True)
        elif data_strategy == 'balance':
            total_len = getattr(cfg, 'total_data_len', 'auto')
            print(f"Using [balance] strategy with total_data_len : {total_len}...")
            trainset_loader = MultipleDatasets(trainset3d_loader + trainset2d_loader + trainset_humandata_loader, 
                                                 make_same_len=True, total_len=total_len, verbose=True)
        else:
            # original strategy implementation
            valid_loader_num = 0
            if len(trainset3d_loader) > 0:
                trainset3d_loader = [MultipleDatasets(trainset3d_loader, make_same_len=False)]
                valid_loader_num += 1
            else:
                trainset3d_loader = []
            if len(trainset2d_loader) > 0:
                trainset2d_loader = [MultipleDatasets(trainset2d_loader, make_same_len=False)]
                valid_loader_num += 1
            else:
                trainset2d_loader = []
            if len(trainset_humandata_loader) > 0:
                trainset_humandata_loader = [MultipleDatasets(trainset_humandata_loader, make_same_len=False)]
                valid_loader_num += 1

            if valid_loader_num > 1:
                trainset_loader = MultipleDatasets(trainset3d_loader + trainset2d_loader + trainset_humandata_loader, make_same_len=True)
            else:
                trainset_loader = MultipleDatasets(trainset3d_loader + trainset2d_loader + trainset_humandata_loader, make_same_len=False)

        self.itr_per_epoch = math.ceil(len(trainset_loader) / cfg.num_gpus / cfg.train_batch_size)

        if self.distributed:
            self.logger_info(f"Total data length {len(trainset_loader)}.")
            rank, world_size = get_dist_info()
            self.logger_info("Using distributed data sampler.")
            
            sampler_train = DistributedSampler(trainset_loader, world_size, rank, shuffle=True)
            self.batch_generator = DataLoader(dataset=trainset_loader, batch_size=cfg.train_batch_size,
                                          shuffle=False, num_workers=cfg.num_thread, sampler=sampler_train,
                                          pin_memory=True, persistent_workers=True if cfg.num_thread > 0 else False, drop_last=True)
        else:
            self.batch_generator = DataLoader(dataset=trainset_loader, batch_size=cfg.num_gpus * cfg.train_batch_size,
                                          shuffle=True, num_workers=cfg.num_thread,
                                          pin_memory=True, drop_last=True)

    def _make_model(self):
        # prepare network
        self.logger_info("Creating graph and optimizer...")
        model = get_model('train')

        if getattr(cfg, 'fine_tune', None) == 'backbone':
            print("Fine-tuning [backbone]...")
            for module in model.head:
                for param in module.parameters():
                    param.requires_grad = False
            for module in model.neck:
                for param in module.parameters():
                    param.requires_grad = False

        elif getattr(cfg, 'fine_tune', None) == 'neck_and_head':
            print("Fine-tuning [neck and head]...")
            for param in model.encoder.parameters():
                param.requires_grad = False
        
        elif getattr(cfg, 'fine_tune', None) == 'head':
            print("Fine-tuning [head]...")
            for param in model.encoder.parameters():
                param.requires_grad = False
            for module in model.neck:
                for param in module.parameters():
                    param.requires_grad = False
        
        
        # ddp
        if self.distributed:
            self.logger_info("Using distributed data parallel.")
            model.cuda()
            if hasattr(cfg, 'syncbn') and cfg.syncbn:
                self.logger_info("Using sync batch norm layers.")

                process_groups = get_process_groups()
                process_group = process_groups[get_group_idx()]
                syncbn_model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model, process_group)
                model = torch.nn.parallel.DistributedDataParallel(
                    syncbn_model, device_ids=[self.gpu_idx],
                    find_unused_parameters=True) 
            else:
                model = torch.nn.parallel.DistributedDataParallel(
                    model, device_ids=[self.gpu_idx],
                    find_unused_parameters=True)

        else:
        # dp
            model = DataParallel(model).cuda()
               
        if self.use_lora:
            lora.mark_only_lora_as_trainable(model.module.encoder, bias='lora_only')  ## lora만 학습 가능하게 수정


        optimizer = self.get_optimizer(model)
        
        if hasattr(cfg, "scheduler"):
            if cfg.scheduler == 'cos':
                scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, cfg.end_epoch * self.itr_per_epoch,
                                                               eta_min=1e-6)
            elif cfg.scheduler == 'step':
                scheduler = torch.optim.lr_scheduler.StepLR(optimizer, cfg.step_size, gamma=cfg.gamma, 
                                                            last_epoch=- 1, verbose=False)                                           

        else:
            scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, cfg.end_epoch * self.itr_per_epoch,
                                                               eta_min=getattr(cfg,'min_lr',1e-6))
        if cfg.continue_train:
            if self.distributed:
                start_epoch, model, optimizer = self.load_model(model, optimizer)
            else:
                start_epoch, model, optimizer = self.load_model(model, optimizer)
        else:
            start_epoch = 0
        model.train()

        self.scheduler = scheduler
        self.start_epoch = start_epoch
        self.model = model
        self.optimizer = optimizer

    def logger_info(self, info):
        if self.distributed:
            if is_main_process():
                self.logger.info(info)
        else:
            self.logger.info(info)


class Tester(Base):
    def __init__(self, test_epoch=None):
        if test_epoch is not None:
            self.test_epoch = int(test_epoch)
        super(Tester, self).__init__(log_name='test_logs.txt')

    def _make_batch_generator(self):
        # data load and construct batch generator
        self.logger.info("Creating dataset...")
        testset_loader = eval(cfg.testset)(transforms.ToTensor(), "test")
        batch_generator = DataLoader(dataset=testset_loader, batch_size=cfg.num_gpus * cfg.test_batch_size,
                                     shuffle=False, num_workers=cfg.num_thread, pin_memory=True)

        self.testset = testset_loader
        self.batch_generator = batch_generator

    def _make_model(self):
        self.logger.info('Load checkpoint from {}'.format(cfg.pretrained_model_path))

        # prepare network
        self.logger.info("Creating graph...")
        model = get_model('test')
        model = DataParallel(model).cuda()
        if not getattr(cfg, 'random_init', False):
            ckpt = torch.load(cfg.pretrained_model_path, map_location=torch.device('cpu'))

            from collections import OrderedDict
            new_state_dict = OrderedDict()
            for k, v in ckpt['network'].items():
                # set_trace()
                if 'module' not in k:
                    k = 'module.' + k
                k = k.replace('backbone', 'encoder').replace('body_rotation_net', 'body_regressor').replace(
                    'hand_rotation_net', 'hand_regressor')
                new_state_dict[k] = v
            self.logger.warning("Attention: Strict=False is set for checkpoint loading. Please check manually.")

            skip_prefixes = ('module.id_head.', 'id_head.')
            filtered_state = OrderedDict(
                (k, v) for k, v in new_state_dict.items()
                if not k.startswith(skip_prefixes)
            )


            model.load_state_dict(filtered_state, strict=False)
            model.eval()
        else:
            print('Random init!!!!!!!')

        self.model = model

    def _evaluate(self, outs, cur_sample_idx):
        eval_result = self.testset.evaluate(outs, cur_sample_idx)
        return eval_result

    def _print_eval_result(self, eval_result):
        self.testset.print_eval_result(eval_result)

class Demoer(Base):
    def __init__(self, test_epoch=None):
        if test_epoch is not None:
            self.test_epoch = int(test_epoch)
        super(Demoer, self).__init__(log_name='test_logs.txt')

    def _make_batch_generator(self, demo_scene):
        # data load and construct batch generator
        self.logger.info("Creating dataset...")
        from data.UBody.UBody import UBody
        testset_loader = UBody(transforms.ToTensor(), "demo", demo_scene) # eval(demoset)(transforms.ToTensor(), "demo")
        batch_generator = DataLoader(dataset=testset_loader, batch_size=cfg.num_gpus * cfg.test_batch_size,
                                     shuffle=False, num_workers=cfg.num_thread, pin_memory=True)

        self.testset = testset_loader
        self.batch_generator = batch_generator

    def _make_model(self):
        self.logger.info('Load checkpoint from {}'.format(cfg.pretrained_model_path))

        # prepare network
        self.logger.info("Creating graph...")
        model = get_model('test')
        model = DataParallel(model).cuda()
        ckpt = torch.load(cfg.pretrained_model_path)

        from collections import OrderedDict
        new_state_dict = OrderedDict()
        for k, v in ckpt['network'].items():
            if 'module' not in k:
                k = 'module.' + k
            k = k.replace('module.backbone', 'module.encoder').replace('body_rotation_net', 'body_regressor').replace(
                'hand_rotation_net', 'hand_regressor')
            new_state_dict[k] = v
        model.load_state_dict(new_state_dict, strict=False)
        model.eval()

        self.model = model

    def _evaluate(self, outs, cur_sample_idx):
        eval_result = self.testset.evaluate(outs, cur_sample_idx)
        return eval_result