QAT issue of yolox

[ianWeichengZhang]

Hello, I am trying to make yolox QAT work. Following the tutorial, I first conducted PTQ quantization on the yolox model, which includes CLE and BC operations. After that, I used aimet_torch.quantsim.QuantizationSimModel.model to obtain a quantized model with the quant_op added, and replaced the original yolox.model in the yolox trainer. However, the forward function now returns a tensor, which is different from the tuple returned by the original yolox.model(). Could you please tell me how should I modify my code?

Error:

2024-01-29 05:45:16 | INFO     | yolox.core.trainer:344 - Training of experiment is done and the best AP is 0.00
2024-01-29 05:45:16 | ERROR    | yolox.core.launch:98 - An error has been caught in function 'launch', process 'MainProcess' (2285414), thread 'MainThread' (139938635503424):
Traceback (most recent call last):

  File "/workspace/workdir_trt/YOLOX/tools/train.py", line 146, in <module>
    launch(
    └ <function launch at 0x7f4406e45dc0>

> File "/workspace/workdir_trt/YOLOX/yolox/core/launch.py", line 98, in launch
    main_func(*args)
    │          └ (╒═══════════════════╤═══════════════════════════════════════════════════════════════════════════════════════════════════════...
    └ <function main at 0x7f4407be6820>

  File "/workspace/workdir_trt/YOLOX/tools/train.py", line 123, in main
    trainer.train()
    │       └ <function Trainer.train at 0x7f4407be6940>
    └ <yolox.core.trainer.Trainer object at 0x7f4407bec910>

  File "/workspace/workdir_trt/YOLOX/yolox/core/trainer.py", line 193, in train
    self.train_in_epoch()
    │    └ <function Trainer.train_in_epoch at 0x7f4407be69d0>
    └ <yolox.core.trainer.Trainer object at 0x7f4407bec910>

  File "/workspace/workdir_trt/YOLOX/yolox/core/trainer.py", line 202, in train_in_epoch
    self.train_in_iter()
    │    └ <function Trainer.train_in_iter at 0x7f4407be6a60>
    └ <yolox.core.trainer.Trainer object at 0x7f4407bec910>

  File "/workspace/workdir_trt/YOLOX/yolox/core/trainer.py", line 208, in train_in_iter
    self.train_one_iter()
    │    └ <function Trainer.train_one_iter at 0x7f4407be6af0>
    └ <yolox.core.trainer.Trainer object at 0x7f4407bec910>

  File "/workspace/workdir_trt/YOLOX/yolox/core/trainer.py", line 229, in train_one_iter
    loss = outputs["total_loss"]
           └ tensor([[[8.0200e+00, 9.7085e+00, 1.6745e+01,  ..., 7.9569e-03,
                       2.4659e-03, 5.6030e-03],
                      [1.2010e+01, 1.41...

IndexError: too many indices for tensor of dimension 3

Here is the code I modified.

in yolox.trainer:
added:

def apply_cross_layer_equalization(model: torch.nn.Module, input_shape: tuple):
    """
        Applying CLE on the model inplace consists of:
            Batch Norm Folding
            Cross Layer Scaling
            High Bias Fold
            Converts any ReLU6 into ReLU.

        :param model: the loaded model
        :param input_shape: the shape of the input to the model
        :return:
    """

    equalize_model(model, input_shape)


def apply_bias_correction(model: torch.nn.Module, data_loader):
    """
        Applies Bias-Correction on the model.
        :param model: The model to quantize
        :param evaluator: Evaluator used during quantization
        :param dataloader: DataLoader used during quantization
        :param logdir: Log directory used for storing log files
        :return: None
    """
        # Rounding mode can be 'nearest' or 'stochastic'
    rounding_mode = 'nearest'

        # Number of samples used during quantization
    num_quant_samples = 16

        # Number of samples used for bias correction
    num_bias_correct_samples = 16

    params = QuantParams(weight_bw=8, act_bw=8, round_mode=rounding_mode, quant_scheme='tf_enhanced')

        # Perform Bias Correction
    bias_correction.correct_bias(model.to(device="cuda"), params, num_quant_samples=num_quant_samples,
                                 data_loader=data_loader, num_bias_correct_samples=num_bias_correct_samples)

def forward_pass(decoder, model, data_loader):
    """forward pass for compute encodings"""
    #pylint:disable = no-member
    tensor_type = torch.cuda.FloatTensor
    model = model.eval()

    for imgs, _, info_imgs, ids in tqdm(data_loader):
        with torch.no_grad():
            imgs = imgs.type(tensor_type)
            outputs = model(imgs)
            if decoder is not None:
                outputs = decoder(outputs, dtype=outputs.type())

def calculate_quantsim_accuracy(model: torch.nn.Module, evaluator: aimet_common.defs.EvalFunction, dataloader,
                                use_cuda: bool = False, logdir: str = '') -> Tuple[torch.nn.Module, float]:
    """
    Calculates model accuracy on quantized simulator and returns quantized model with accuracy.

    :param model: the loaded model
    :param evaluator: the Eval function to use for evaluation
    :param iterations: No of batches to use in computing encodings.
                       Not used in image net dataset
    :param num_val_samples_per_class: No of samples to use from every class in
                                      computing encodings. Not used in pascal voc
                                      dataset
    :param use_cuda: the cuda device.
    :return: a tuple of quantsim and accuracy of model on this quantsim
    """

    input_shape = (1, 3, 640, 640)
    if use_cuda:
        model.to(torch.device('cuda'))
        dummy_input = torch.rand(input_shape).cuda()
    else:
        dummy_input = torch.rand(input_shape)

    # apply_cross_layer_equalization(model,input_shape)
    # apply_bias_correction(model,dataloader)

    # Number of batches to use for computing encodings
    # Only 5 batches are used here to speed up the process, also the
    # number of images in these 5 batches should be sufficient for
    # compute encodings
    iterations = 5

    net_dataloader = get_data_loader(
        dataset_path='/workspace/datasets',
        img_size=(640,640),
        batch_size=64,
        num_workers=4,
    )

    quantsim = QuantizationSimModel(model=model, quant_scheme='tf_enhanced',
                                    dummy_input=dummy_input, rounding_mode='nearest',
                                    default_output_bw=8, default_param_bw=8, in_place=False)

    forward_func = partial(forward_pass, None)

    quantsim.compute_encodings(forward_func, forward_pass_callback_args=net_dataloader)

    #accuracy = evaluator.evaluate(quantsim.model)

    return quantsim, accuracy

modified

def before_train(self):
        logger.info("args: {}".format(self.args))
        logger.info("exp value:\n{}".format(self.exp))

        # model related init
        torch.cuda.set_device(self.local_rank)
        model = self.exp.get_model()
        logger.info(
            "Model Summary: {}".format(get_model_info(model, self.exp.test_size))
        )
        model.to(self.device)

        # solver related init
        self.optimizer = self.exp.get_optimizer(self.args.batch_size)

        # value of epoch will be set in `resume_train`
        model = self.resume_train(model)


        self.evaluator = self.exp.get_evaluator(
            batch_size=self.args.batch_size, is_distributed=self.is_distributed
        )
        
        # data related init
        self.no_aug = self.start_epoch >= self.max_epoch - self.exp.no_aug_epochs
        self.train_loader = self.exp.get_data_loader(
            batch_size=self.args.batch_size,
            is_distributed=self.is_distributed,
            no_aug=self.no_aug,
            cache_img=self.args.cache,
        )
        logger.info("init prefetcher, this might take one minute or less...")
        self.prefetcher = DataPrefetcher(self.train_loader)
        # max_iter means iters per epoch
        self.max_iter = len(self.train_loader)
        # logger.info("PTQ complete, map50 accuracy:{}".format(map50))
        
        quantsim, map50 = calculate_quantsim_accuracy(model=model, 
                                                  evaluator=self.evaluator,
                                                  dataloader=self.train_loader,
                                                  use_cuda=True,
                                                  logdir=''
                                                  )

        model = quantsim.model

        

        self.lr_scheduler = self.exp.get_lr_scheduler(
            self.exp.basic_lr_per_img * self.args.batch_size, self.max_iter
        )
        if self.args.occupy:
            occupy_mem(self.local_rank)

        if self.is_distributed:
            model = DDP(model, device_ids=[self.local_rank], broadcast_buffers=False)

        if self.use_model_ema:
            self.ema_model = ModelEMA(model, 0.9998)
            self.ema_model.updates = self.max_iter * self.start_epoch 
        
        
        
        

        self.model = model

        
        # Tensorboard and Wandb loggers
        if self.rank == 0:
            if self.args.logger == "tensorboard":
                self.tblogger = SummaryWriter(os.path.join(self.file_name, "tensorboard"))
            elif self.args.logger == "wandb":
                self.wandb_logger = WandbLogger.initialize_wandb_logger(
                    self.args,
                    self.exp,
                    self.evaluator.dataloader.dataset
                )
            else:
                raise ValueError("logger must be either 'tensorboard' or 'wandb'")

        logger.info("Training start...")
        logger.info("\n{}".format(model))