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:
- 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.
- 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