code/cube3d/training/bert_infer.py

"""
Simple training loop; Boilerplate that could apply to any arbitrary neural network,
so nothing in this file really has anything to do with GPT specifically.
"""
from typing import Optional, Tuple, List
import time
import os
from collections import defaultdict

from accelerate import Accelerator
import torch
from torch.nn import functional as F
from torch.utils.data.dataloader import DataLoader
# from mingpt.utils import CfgNode as CN
from omegaconf import DictConfig as CN  # 使用omegaconf作为配置替代
from cube3d.training.utils import save_model_weights, mask_cross_entropy, normalize_bboxs, top_k_prob_mask, visualize_token_probabilities, visualize_max_prob_distribution
from cube3d.training.process_single_ldr import logits2ldr, logits2ldrot, logits2ldrp, logits2flatldrp, logits2flatldrpr, logits2botldrpr
from cube3d.inference.utils import load_model_weights          
from tqdm import tqdm


def generate_tokens(
    engine,
    prompt,
    inputs_ids,
    latent,
    resolution_base=8.0,
    disable_postprocess=False,
    top_p=None,
    bounding_box_xyz=None,
    strategy=None,
    mode=None
):
    output_ids = engine.t2t(
        #[prompt],
        prompt,
        #use_kv_cache=True,
        inputs_ids=inputs_ids,
        latent=latent,
        use_kv_cache=False,
        resolution_base=resolution_base,
        top_p=top_p,
        bounding_box_xyz=bounding_box_xyz,
        strategy=strategy,
        mode=mode
    )

    return output_ids

class Infer:

    @staticmethod
    def get_default_config():
        C = CN()
        # device to train on
        C.device = 'auto'
        # dataloder parameters
        C.num_workers = 4
        # optimizer parameters
        C.max_iters = None
        C.batch_size = 4
        C.learning_rate = 3e-4
        C.betas = (0.9, 0.95)
        C.weight_decay = 0.1 # only applied on matmul weights
        C.grad_norm_clip = 1.0
        C.save_interval = None
        return C

    def __init__(
        self, 
        config, 
        engine, 
        train_dataset,
        accelerator,
        tb,
        prompt: str,
        indices: Optional[List[int]] = None,
        resolution_base: float = 8.0,
        disable_postprocessing: bool = False,
        top_p: float = None,
        bounding_box_xyz: Optional[Tuple[float]] = None,
        save_gpt_ckpt_path: str = None,
        mode: str = 'train'
        ):
        self.config = config
        self.engine = engine
        self.model = engine.gpt_model
        self.optimizer = None
        self.callbacks = defaultdict(list)
        self.train_dataset = train_dataset
        self.accelerator = accelerator

        # Training parameters
        self.prompt = prompt
        self.targets = indices
        self.resolution_base = resolution_base
        self.disable_postprocessing = disable_postprocessing
        self.top_p = top_p
        self.bounding_box_xyz = bounding_box_xyz
        self.save_gpt_ckpt_path = save_gpt_ckpt_path

        # determine the device we'll train on
        if config.device == 'auto':
            self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        else:
            self.device = config.device

        self.model = self.model.to(self.device)
        print("running on device", self.device)

        # variables that will be assigned to trainer class later for logging and etc
        self.iter_num = 0
        self.iter_time = 0.0
        self.iter_dt = 0.0

        self.tb_writer = tb
        self.mode = mode

        self.probs = []


    def add_callback(self, onevent: str, callback):
        self.callbacks[onevent].append(callback)

    def set_callback(self, onevent: str, callback):
        self.callbacks[onevent] = [callback]

    def trigger_callbacks(self, onevent: str):
        for callback in self.callbacks.get(onevent, []):
            callback(self)

    def run(self):
        model, config = self.model, self.config
        # setup the optimizer
        #self.optimizer = self.engine.configure_optimizers(config)
        self.optimizer, self.scheduler = self.engine.configure_optimizers_scratch_linear(config) #self.engine.configure_optimizers_lora_linear(config)

        # setup the dataloader
        train_loader = DataLoader(
            self.train_dataset,
            shuffle=False if self.mode!='train' else True,
            batch_size=config.batch_size,
        )

        model.train()

        model, self.optimizer, train_loader = self.accelerator.prepare(model, self.optimizer, train_loader)

        self.iter_num = 0
        self.iter_time = time.time()
        data_iter = iter(train_loader)
        ema_loss_for_log = 0.0
        ema_ploss_for_log = 0.0
        ema_rloss_for_log = 0.0
        ema_dloss_for_log = 0.0
        ema_floss_for_log = 0.0
        
        #loss
        dat_num = 1217 #301
        x_num = 251 #213
        y_num = 215 #73
        z_num = 525 #411
        rot_num = 24
        shift = 0
        stride = 5
        attr_shift = stride-3 #with dat and rot，+1 for bert
        bert_shift = 1

        x = x_num
        xy = x_num + y_num + rot_num
        xyz = x_num + y_num + z_num + rot_num

        progress_bar = tqdm(range(0, config.max_iters), desc="Training progress")
        #while True:
        for self.iter_num in range(0, config.max_iters+1): 
            # fetch the next batch (x, y) and re-init iterator if needed
            try:
                batch = next(data_iter)
            except StopIteration:
                data_iter = iter(train_loader)
                batch = next(data_iter)

            #batch = [t['latent'].to(self.device) for t in batch]
            self.prompt, self.targets, self.box = batch['prompt'], batch['target'].to(self.device), batch['bbox']
            #self.targets = batch['latent'].to(self.device)
            targets = self.targets.clone()
#             import ipdb; ipdb.set_trace()
            logits, inputs_ids, strategy, mask, cut_idx = generate_tokens(
                self.engine,
                self.prompt,
                targets,
                None,
                self.resolution_base,
                self.disable_postprocessing,
                self.top_p,
                #self.bounding_box_xyz,
                normalize_bboxs(self.box.float(), [x_num-1, y_num-1, z_num-1]), #batch_normalization(self.box)
                1,
                self.mode
            )

            
            # rotation_loss = F.cross_entropy(
            #     logits[:,:-1,:rot_num].permute(0, 2, 1),
            #     inputs_ids[:,shift:,:rot_num].argmax(-1), 
            # )

            
            # px_loss = mask_cross_entropy(rot_num, x+rot_num, self.box[:, 0], logits, inputs_ids, shift)
            # py_loss = mask_cross_entropy(x+rot_num, xy, self.box[:, 1], logits, inputs_ids, shift)
            # pz_loss = mask_cross_entropy(xy, xyz, self.box[:, 2], logits, inputs_ids, shift)

            px_loss =  F.cross_entropy(
                logits[:,1+attr_shift+bert_shift:-1:stride,rot_num+1:x+rot_num+1+1].permute(0, 2, 1), 
                inputs_ids[:,shift:,-5], 
                ignore_index=-1 #+1 for padding
            )
            py_loss =  F.cross_entropy(
                logits[:,0+attr_shift+bert_shift:-2:stride,x+rot_num+2:xy+3].permute(0, 2, 1), 
                inputs_ids[:,shift:,-4], 
                ignore_index=-1
            )
            pz_loss =  F.cross_entropy(
                logits[:,2+attr_shift+bert_shift::stride,xy+3:xyz+4].permute(0, 2, 1), 
                inputs_ids[:,shift:,-3], 
                ignore_index=-1
            )

            position_loss = px_loss + py_loss + pz_loss
            
            # dat_loss =  F.cross_entropy(
            #     logits[:,0:-4:stride,:dat_num+1].permute(0, 2, 1), 
            #     inputs_ids[:,shift:,-6], 
            #     ignore_index=-1
            # )

            rotation_loss =  F.cross_entropy(
                logits[:,1+bert_shift:-3:stride,:rot_num+1].permute(0, 2, 1), 
                inputs_ids[:,shift:,-7], 
                ignore_index=-1
            )


            # flag_loss = F.cross_entropy(
            #     logits[:,:-1,xyz+dat_num:xyz+dat_num+2].permute(0, 2, 1), 
            #     inputs_ids[:,shift:,xyz+dat_num:xyz+dat_num+2].argmax(-1), 
            # )

            # flag_loss = F.cross_entropy(
            #     logits[:,:-1,-2:].permute(0, 2, 1), 
            #     inputs_ids[:,shift:,-2:].argmax(-1), 
            # )

            lambda_posiition = 1.0
            lambda_rotation = 1.0
            lambda_dat = 1.0
            lambda_flag = 50.0

            self.loss = lambda_posiition * position_loss #+ \
                        #lambda_rotation * rotation_loss #+ \
                        #lambda_flag * flag_loss
                        #lambda_dat * dat_loss + \

            if strategy==1 or strategy==2:
                self.loss+=lambda_rotation * rotation_loss


            targets = self.targets.clone()
            # mask_topk, mask_inv = top_k_prob_mask(F.softmax(logits[:, 1:-3:stride, :rot_num+1], dim=2), cut_idx, top_percent=0.5)
            # targets[:,shift:,-7][mask_topk] = logits[:,1:-3:stride,:rot_num+1].permute(0, 2, 1).argmax(dim=1)[mask_topk]
            # targets[:,shift:,-7][mask_inv] = self.engine.gpt_model.rot_num+1

            targets[:,shift:,-7] = logits[:,1:-3:stride,:rot_num+1].permute(0, 2, 1).argmax(dim=1)
            #targets[:,shift:,-4] = logits_y[:,0+attr_shift:-2:stride,x+rot_num+2:xy+3].permute(0, 2, 1).argmax(dim=1)
            # fig = visualize_token_probabilities(
            #     probs=F.softmax(logits[:, 1:-3:stride, :rot_num+1], dim=2),
            #     cut_idx=cut_idx,
            #     sample_idx=0,
            #     tokens_per_page=10,
            #     save_dir='token_probability_pages'  # 图片会保存到这个文件夹
            # )

            logits_x, inputs_ids, strategy, mask, cut_idx = generate_tokens(
                self.engine,
                self.prompt,
                targets,
                None,
                self.resolution_base,
                self.disable_postprocessing,
                self.top_p,
                #self.bounding_box_xyz,
                normalize_bboxs(self.box.float(), [x_num-1, y_num-1, z_num-1]), #batch_normalization(self.box)
                0,
                self.mode
            )

            logits_x[:,1+bert_shift:-3:stride,:rot_num+1] = logits[:,1+bert_shift:-3:stride,:rot_num+1]
            output_dir = "train_d2r2p_scratch_whole_perm10re"
            os.makedirs(output_dir, exist_ok=True)
            logits2botldrpr(logits_x[0].cpu().detach().numpy(), inputs_ids[0].cpu().detach().numpy(), stride, 0, output_file=os.path.join(output_dir, f"test_d2r2p_{self.iter_num}_scratch_0p85_bert.ldr"))


            #targets = self.targets.clone()
            #targets[:,shift:,-7] = logits[:,1+bert_shift:-3:stride,:rot_num+1].permute(0, 2, 1).argmax(dim=1)

            #mask_x, mask_x_inv = top_k_prob_mask(F.softmax(logits_x[:,1+attr_shift+bert_shift:-1:stride,rot_num+1:x+rot_num+1+1], dim=2), cut_idx, top_percent=0.3)
            #mask_y, mask_y_inv = top_k_prob_mask(F.softmax(logits_x[:,0+attr_shift+bert_shift:-2:stride,x+rot_num+2:xy+3], dim=2), cut_idx, top_percent=0.3)
            #mask_z, mask_z_inv = top_k_prob_mask(F.softmax(logits_x[:,2+attr_shift+bert_shift::stride,xy+3:xyz+4], dim=2), cut_idx, top_percent=0.3)

            #targets[:,shift:,-5][mask_x] = logits_x[:,1+attr_shift+bert_shift:-1:stride,rot_num+1:x+rot_num+1+1].permute(0, 2, 1).argmax(dim=1)[mask_x]
            #targets[:,shift:,-5][mask_x_inv] = self.engine.gpt_model.x_num+1
            #targets[:,shift:,-4][mask_y] = logits_x[:,0+attr_shift+bert_shift:-2:stride,x+rot_num+2:xy+3].permute(0, 2, 1).argmax(dim=1)[mask_y]
            #targets[:,shift:,-4][mask_y_inv] = self.engine.gpt_model.y_num+1
            #targets[:,shift:,-3][mask_z] = logits_x[:,2+attr_shift+bert_shift::stride,xy+3:xyz+4].permute(0, 2, 1).argmax(dim=1)[mask_z]
            #targets[:,shift:,-3][mask_z_inv] = self.engine.gpt_model.z_num+1

            # fig = visualize_token_probabilities(
            #     probs=F.softmax(logits_x[:,1+attr_shift+bert_shift:-1:stride,rot_num+1:x+rot_num+1+1], dim=2),
            #     cut_idx=cut_idx,
            #     sample_idx=0,
            #     tokens_per_page=10,
            #     save_dir='token_probability_pages_x'  # 图片会保存到这个文件夹
            # )

            # fig = visualize_token_probabilities(
            #     probs=F.softmax(logits_x[:,0+attr_shift+bert_shift:-2:stride,x+rot_num+2:xy+3], dim=2),
            #     cut_idx=cut_idx,
            #     sample_idx=0,
            #     tokens_per_page=10,
            #     save_dir='token_probability_pages_y'  # 图片会保存到这个文件夹
            # )
        
            # fig = visualize_token_probabilities(
            #     probs=F.softmax(logits_x[:,2+attr_shift+bert_shift::stride,xy+3:xyz+4], dim=2),
            #     cut_idx=cut_idx,
            #     sample_idx=0,
            #     tokens_per_page=10,
            #     save_dir='token_probability_pages_z'  # 图片会保存到这个文件夹
            # )
            # fig = visualize_max_prob_distribution(
            #     probs=F.softmax(logits_x[:,2+attr_shift+bert_shift::stride,xy+3:xyz+4], dim=2),
            #     cut_idx=cut_idx,
            #     sample_idx=0,
            #     bins=20,  # 0-1分成20个区间（每个区间0.05）
            #     figsize=(12, 6)
            # )
            # fig.savefig(f'token_max_probability_distribution_iter{self.iter_num}.png')
            # current_probs = torch.softmax(logits[:,1+bert_shift:-3:stride,:rot_num+1], dim=2)
            # self.probs.append(current_probs[:, :min(int(cut_idx[0]), current_probs.shape[1]), :])
            # logits_p, inputs_ids, strategy, mask, cut_idx = generate_tokens(
            #     self.engine,
            #     self.prompt,
            #     targets,
            #     None,
            #     self.resolution_base,
            #     self.disable_postprocessing,
            #     self.top_p,
            #     #self.bounding_box_xyz,
            #     normalize_bboxs(self.box.float(), [x_num-1, y_num-1, z_num-1]), #batch_normalization(self.box)
            #     0
            # )

            # logits_p[:,1+bert_shift:-3:stride,:rot_num+1] = logits[:,1+bert_shift:-3:stride,:rot_num+1]
            # logits2botldrpr(logits_p[0].cpu().detach().numpy(), inputs_ids[0].cpu().detach().numpy(), stride, 0, output_file=f"gt_d2r2p2p_scratch_0p85_bot/test_d2r2p2p_{self.iter_num}_scratch_0p85_bert.ldr")

            # targets = self.targets.clone()
            # targets[:,shift:,-7] = logits[:,1:-3:stride,:rot_num+1].permute(0, 2, 1).argmax(dim=1)
            # targets[:,shift:,-4] = logits_y[:,0+attr_shift:-2:stride,x+rot_num+2:xy+3].permute(0, 2, 1).argmax(dim=1)
            # targets[:,shift:,-5] = logits_x[:,1+attr_shift:-1:stride,rot_num+1:x+rot_num+1+1].permute(0, 2, 1).argmax(dim=1)
            # logits_z, inputs_ids, strategy = generate_tokens(
            #     self.engine,
            #     self.prompt,
            #     targets,
            #     None,
            #     self.resolution_base,
            #     self.disable_postprocessing,
            #     self.top_p,
            #     #self.bounding_box_xyz,
            #     normalize_bboxs(self.box.float(), [x_num-1, y_num-1, z_num-1]), #batch_normalization(self.box)
            #     3
            # )

            # backprop and update the parameters
            model.zero_grad(set_to_none=True)
            # #if self.mode!='train':
            # logits_z[:,1:-3:stride,:rot_num+1] = logits[:,1:-3:stride,:rot_num+1]
            # logits_z[:,0+attr_shift:-2:stride,x+rot_num+2:xy+3] = logits_y[:,0+attr_shift:-2:stride,x+rot_num+2:xy+3]
            # logits_z[:,1+attr_shift:-1:stride,rot_num+1:x+rot_num+1+1] = logits_x[:,1+attr_shift:-1:stride,rot_num+1:x+rot_num+1+1]

            if self.iter_num>4:
                break

            # self.accelerator.backward(self.loss)
            # torch.nn.utils.clip_grad_norm_(model.parameters(), config.grad_norm_clip)
            # self.optimizer.step()
            # self.scheduler.step()

            with torch.no_grad():
                # Progress bar
                ema_loss_for_log = 0.4 * self.loss.item() + 0.6 * ema_loss_for_log
                ema_ploss_for_log = 0.4 * position_loss.item() + 0.6 * ema_ploss_for_log
                ema_rloss_for_log = 0.4 * rotation_loss.item() + 0.6 * ema_rloss_for_log
                #ema_dloss_for_log = 0.4 * dat_loss.item() + 0.6 * ema_dloss_for_log
                #ema_floss_for_log = 0.4 * flag_loss.item() + 0.6 * ema_floss_for_log
                if self.iter_num % 10 == 0:
                    progress_bar.set_postfix({"Loss": f"{ema_loss_for_log:.{7}f}",
                                            "Positon_Loss": f"{ema_ploss_for_log:.{7}f}",
                                             "Rotation_Loss": f"{ema_rloss_for_log:.{7}f}",
                                             #"Dat_Loss": f"{ema_dloss_for_log:.{7}f}",
                                             #"Flag_Loss": f"{ema_floss_for_log:.{7}f}",
                                             })
                    progress_bar.update(10)
                
                #logits2ldr(logits[0].cpu().detach().numpy())
                
                if  (self.iter_num % config.save_interval == 0 and self.iter_num != 0):
                    if self.accelerator.is_main_process:
                        save_model_weights(
                        self.engine.gpt_model,
                        self.save_gpt_ckpt_path,
                        )             
                        # self.engine.gpt_model.save_pretrained(self.save_gpt_ckpt_path)
                        # torch.save({
                        #     "ldr_proj": self.engine.gpt_model.ldr_proj.state_dict(),
                        #     "ldr_head": self.engine.gpt_model.ldr_head.state_dict(),
                        #     "rte": self.engine.gpt_model.rte.state_dict(),
                        #     "dte": self.engine.gpt_model.dte.state_dict(),
                        #     "xte": self.engine.gpt_model.xte.state_dict(),
                        #     "yte": self.engine.gpt_model.yte.state_dict(),
                        #     "zte": self.engine.gpt_model.zte.state_dict(),
                        # }, f"{self.save_gpt_ckpt_path}/unfrozen_weights.pth")


                if self.tb_writer: #and self.accelerator.is_main_process:
                    self.tb_writer.add_scalar(f'train_loss/position_loss', position_loss.item(), self.iter_num)
                    self.tb_writer.add_scalar(f'train_loss/rotation_loss', rotation_loss.item(), self.iter_num)
                    #self.tb_writer.add_scalar(f'train_loss/dat_loss', dat_loss.item(), self.iter_num)
                    #self.tb_writer.add_scalar(f'train_loss/flag_loss', flag_loss.item(), self.iter_num)
                    self.tb_writer.add_scalar(f'train_loss/total_loss', self.loss.item(), self.iter_num)

                if self.iter_num == config.max_iters:
                    progress_bar.close()