code/cube3d/training/trainer.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 cube3d.training.utils import save_model_weights, mask_cross_entropy, normalize_bboxs, top_k_prob_mask
from cube3d.training.process_single_ldr import logits2ldr, logits2ldrot, logits2ldrp, logits2flatldrp, logits2flatldrpr
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
):
    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
    )

    return output_ids

class Trainer:

    @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


    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 #286
        x_num = 251
        y_num = 215
        z_num = 525
        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()
            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)
                None
            )

            
            # 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)
            # 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
            # )

            # targets = self.targets.clone()
            # targets[:,shift:,-7] = logits_x[:,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[:,1+attr_shift+bert_shift:-1:stride,rot_num+1:x+rot_num+1+1], dim=2), cut_idx, top_percent=0.5)
            # mask_y, mask_y_inv = top_k_prob_mask(F.softmax(logits[:,0+attr_shift+bert_shift:-2:stride,x+rot_num+2:xy+3], dim=2), cut_idx, top_percent=0.5)
            # mask_z, mask_z_inv = top_k_prob_mask(F.softmax(logits[:,2+attr_shift+bert_shift::stride,xy+3:xyz+4], dim=2), cut_idx, top_percent=0.5)

            # 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
            # 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)
            #     None
            # )

            # logits_p[:,1+bert_shift:-3:stride,:rot_num+1] = logits[:,1+bert_shift:-3:stride,:rot_num+1]
            # logits2flatldrpr(logits_p[0].cpu().detach().numpy(), inputs_ids[0].cpu().detach().numpy(), stride, 0, output_file=f"test_rightd2r2p2p_{self.iter_num}_scratch_0p5_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()