Quantizing the Model - 2.0 English

Vitis AI User Guide (UG1414)

Document ID
UG1414
Release Date
2022-01-20
Version
2.0 English
vai_q_pytorch provides a decorator to register an operation or a group of operations as a custom operation which is unknown to XIR.

# Decorator API
def register_custom_op(op_type: str, attrs_list: Optional[List[str]] = None):
  """The decorator is used to register the function as a custom operation.
  Args:
  op_type(str) - the operator type registered into quantizer. 
  The type should not conflict with pytorch_nndct

  attrs_list(Optional[List[str]], optional) - 
  the name list of attributes that define operation flavor. 
  For example, Convolution operation has such attributes as padding, dilation, stride and groups. 
  The order of name in attrs_list should be consistent with that of the arguments list. 
  Default: None

"""

To quantize a model with custom op, two steps are requires to edit the code:

  1. Move the target code into a function and change its calling accordingly. To Pointpillar model, replace the PointPillarsScatter model with a PPScatterV2 function. Check related code in code/test/models/voxelnet.py file.
  2. Decorate this function with decorator API:
    from pytorch_nndct.utils import register_custom_op
    ...
    
    @register_custom_op("PPScatterV2", attrs_list=['ny', 'nx', 'nchannels'])
    def PPScatterV2(ctx, voxel_features, coords, ny, nx, nchannels):
        '''
        input:
        voxel_features: B x 64 x 12000 x 1
        coords: B x 12000 x 4, 4 channels: [batch_idx, z_idx, y_idx, x_idx]
        '''
        batch_size = voxel_features.shape[0]
        # batch_canvas will be the final output.
        batch_canvas = []
    
        for b_idx in range(batch_size):
            # Create the canvas for this sample
            canvas = torch.zeros(nchannels, nx * ny, dtype=voxel_features.dtype,
                                 device=voxel_features.device)
            # Only include non-empty pillars
    
            batch_mask = coords[b_idx, :, 0] > -1
            this_coords = coords[b_idx, batch_mask, :]
            indices = this_coords[:, 2] * nx + this_coords[:, 3]
            indices = indices.type(torch.long)
    
            voxels = voxel_features[b_idx, :, batch_mask, 0]
    
            # Now scatter the blob back to the canvas.
            canvas[:, indices] = voxels
            # Append to a list for later stacking.
            batch_canvas.append(canvas)
    
        # Stack to 3-dim tensor (batch-size, nchannels, nrows*ncols)
        batch_canvas = torch.stack(batch_canvas, 0)
        # Undo the column stacking to final 4-dim tensor
        batch_canvas = batch_canvas.view(batch_size, nchannels, ny, nx)
        return batch_canvas
    
After the target custom op code has been prepared and decorated, add general vai_q_pytorch API functions (check the related code in code/test/test.py)

if quant_mode != 'float':
    max_voxel_num = config.eval_input_reader.max_number_of_voxels
    max_point_num_per_voxel = model_cfg.voxel_generator.max_number_of_points_per_voxel
    aug_voxels = torch.randn((1, 4, max_voxel_num, max_point_num_per_voxel)).to(device)
    # coors = torch.randn((max_voxel_num, 4)).to(device)
    coors = torch.randn((1, max_voxel_num, 4)).to(device)
    quantizer = torch_quantizer(quant_mode=quant_mode,
                                module=net,
                                input_args=(aug_voxels, coors),
                                output_dir=output_dir,
                                device=device,
                                )
    net = quantizer.quant_model
...
...
for example in iter(eval_dataloader):
...
    if quant_mode == 'test' and args.dump_xmodel:
        quantizer.export_xmodel(output_dir=output_dir, deploy_check=True)
        sys.exit()
...
...
if quant_mode == 'calib':
    quantizer.export_quant_config()
After all changes are ready, run script code/test/run_quant.sh to get quantization result files, including xmodel file to compiler (./quantized/VoxelNet_int.xmodel):
sh ./code/test/run_quant.sh