量化模型 - 3.5 简体中文

vai_q_pytorch 提供了一个修饰器，可将一项或一组运算寄存为对 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 names in attrs_list should be consistent with that of the arguments list. 
  Default: None

"""

要使用自定义运算符来量化模型，请遵循以下步骤编辑代码：

1. 将目标代码移入函数中，并对其调用进行相应更改。对于 Pointpillar 模型，请将 PointPillarsScatter 模型替换为 PPScatterV2 函数。检查 the code/test/models/voxelnet.py 文件中的相关代码。
2. 利用修饰器 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 sthe the tacking to final 4-dim tensor
       batch_canvas = batch_canvas.view(batch_size, nchannels, ny, nx)
       return batch_canvas
   

完成目标自定义运算符代码的准备和修饰工作后，添加常规的 vai_q_pytorch API 函数（检查 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() 

准备好所有更改后，运行 code/test/run_quant.sh 脚本以获取量化结果文件，包括编译器的 XMODEL 文件 (./quantized/VoxelNet_int.xmodel)：

sh ./code/test/run_quant.sh