export onnx error

[xiongda777]

when i export fastpillar onnx,torch_scatter.scatter_max is not support by onnx.Do you have any solution?

[StiphyJay]

Due to the MAPE is very slight, you can only export onnx file consists of backbone + neck +head, while the MAPE module is retained in Python style. Besides, you may also refer to the scatter handling of PointPillars in https://github.com/CarkusL/CenterPoint.

[xiongda777]

In fact, I am exporting backbone+neck+head, but the scatter_max operator in pillar_comodules. py is not supported by onnx, and it is not used in pointpillar. this is my export code.Thank you

from det3d.models import build_detector
from det3d.torchie.trainer import get_dist_info, load_checkpoint
from det3d.torchie import Config
from det3d.datasets import build_dataloader, build_dataset
import torch
import torch.nn as nn
from torch.onnx import symbolic_helper


class Model(nn.Module):
    def __init__(self,model):
        super().__init__()
        self.model=model

    def forward(self, xyz,xyz_batch_cnt,pt_features):
        # batch_size = 1

        # data = dict(
        #     points=example['points'],
        #     batch_size=batch_size,
        # )

        # data = self.model.reader(data)

        pillar_features = self.model.backbone(
           xyz,xyz_batch_cnt,pt_features
        )


        x = self.model.neck(pillar_features)

        preds, _ = self.model.bbox_head(x, pillar_features,self.model.test_cfg)
        example={}
        boxes=self.model.bbox_head.predict(example, preds, self.model.test_cfg)
        return boxes

@symbolic_helper.parse_args("v", "v", "i", "v", "i", "i")
def symbolic_scatter_max(g, src, index, dim=-1, out=None, dim_size=None, fill_value=None):
    return (index ,g.op("ScatterElements", out, index, src, axis_i=dim, reduction_s="max"))

torch.onnx.register_custom_op_symbolic("torch_scatter::scatter_max", symbolic_scatter_max, 17)

config="configs/nusc/fastpillars/nusc_fastpillars_015pillar_res34_2dconv_rep_6631.py"
checkpoint="/data2/xd/FastPillars/work_dirs/nusc_fastpillars_015pillar_res34_2dconv_rep_6631/epoch_30.pth"
cfg = Config.fromfile(config)

model = build_detector(cfg.model, train_cfg=None, test_cfg=cfg.test_cfg)

checkpoint = load_checkpoint(model, checkpoint, map_location="cpu")


dataset = build_dataset(cfg.data.test)
data_loader = build_dataloader(
        dataset,
        batch_size=1,
        workers_per_gpu=cfg.data.workers_per_gpu,
        dist=False,
        shuffle=False,
    )


model=model.cuda()

model=Model(model=model)
for i, data_batch in enumerate(data_loader):
    data_batch['metadata'][0]['image_prefix']='data/rcs_data_1203'
    data_batch['shape']=torch.from_numpy(data_batch['shape'])
    data_batch['points'][0]=data_batch['points'][0].cuda()
    data_batch['voxels']=data_batch['voxels'].cuda()
    data_batch['num_points']=data_batch['num_points'].cuda()
    data_batch['num_voxels']=data_batch['num_voxels'].cuda()
    data_batch['coordinates']=data_batch['coordinates'].cuda()
    xyz=data_batch['points'][0][:,:3].cuda().contiguous()
    xyz_batch_cnt=torch.tensor([data_batch['points'][0].shape[0]],dtype=torch.int).cuda().contiguous()
    pt_features=data_batch['points'][0].cuda().contiguous()
    input_names=['xyz','xyz_batch_cnt','pt_features']

    xyz=torch.zeros((640000,3)).cuda().contiguous()
    xyz_batch_cnt=torch.tensor([640000],dtype=torch.int).cuda().contiguous()
    pt_features=torch.zeros((640000,5)).cuda().contiguous()
    torch.onnx.export(model,
                      (xyz,xyz_batch_cnt,pt_features),
                      "test.onnx",
                      do_constant_folding=False,
                      opset_version=17,
                      verbose=True,
                      input_names=input_names,
    )   
    break