Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Convert to ONNX #22

Open
RKGalaxy-Luo opened this issue Aug 23, 2024 · 0 comments
Open

Convert to ONNX #22

RKGalaxy-Luo opened this issue Aug 23, 2024 · 0 comments

Comments

@RKGalaxy-Luo
Copy link

RKGalaxy-Luo commented Aug 23, 2024

Thanks to the author for your wonderful work. I'm going to deploy this model as a part of our project.

I emend the return value in the forward function as followed. (Because the type of SE3 can not be accecpt by ONNX output).

def forward(self, image1, image2, depth1, depth2, intrinsics, iters=12, train_mode=False):

        # Adjust the image and depth input tensor
        image1, image2, depth1, depth2 = prepare_images_and_depths(image1, image2, depth1, depth2)
        Ts, coords0 = self.initializer(image1)
        corr_fn, net, inp = self.features_and_correlation(image1, image2)

        # intrinsics and depth at 1/8 resolution
        intrinsics_r8 = intrinsics / 8.0
        depth1_r8 = depth1[:,3::8,3::8]
        depth2_r8 = depth2[:,3::8,3::8]

        flow_est_list = []
        flow_rev_list = []

        for itr in range(iters):
            Ts = Ts.detach()

            coords1_xyz, _ = pops.projective_transform(Ts, depth1_r8, intrinsics_r8)
            
            coords1, zinv_proj = coords1_xyz.split([2,1], dim=-1)
            zinv, _ = depth_sampler(1.0/depth2_r8, coords1)

            corr = corr_fn(coords1.permute(0,3,1,2).contiguous())
            flow = coords1 - coords0

            dz = zinv.unsqueeze(-1) - zinv_proj
            twist = Ts.log()

            net, mask, ae, delta, weight = \
                self.update_block(net, inp, corr, flow, dz, twist)

            target = coords1_xyz.permute(0,3,1,2) + delta
            target = target.contiguous()

            # Gauss-Newton step
            # Ts = se3_field.step(Ts, ae, target, weight, depth1_r8, intrinsics_r8)
            Ts = se3_field.step_inplace(Ts, ae, target, weight, depth1_r8, intrinsics_r8)

            if train_mode:
                flow2d_rev = target.permute(0,2,3,1)[...,:2] - coords0
                flow2d_rev = se3_field.cvx_upsample(8 * flow2d_rev, mask)

                Ts_up = se3_field.upsample_se3(Ts, mask)
                flow2d_est, flow3d_est, valid = pops.induced_flow(Ts_up, depth1, intrinsics)

                flow_est_list.append(flow2d_est)
                flow_rev_list.append(flow2d_rev)

        if train_mode:
            return flow_est_list, flow_rev_list

        Ts_up = se3_field.upsample_se3(Ts, mask)

        # compute 2d and 3d from from SE3 field (Ts)
        flow2d, flow3d, _ = pops.induced_flow(Ts_up, depth1, intrinsics)
        # undo depth scaling
        flow3d = flow3d / DEPTH_SCALE
        return flow3d

And my Convert2OnnxModel Function is followed as:

def Convert2OnnxModel() -> None:
    print(torch.__version__)
    print(torch.cuda.is_available())

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    print("Begin to Convert Net Model to Onnx Model")
    model_path = '/home/h3c/LuoJiaxuan/RAFT3D/pretrained/raft3d.pth'

    parser = argparse.ArgumentParser()
    parser.add_argument('--model', default='pretrained/raft3d.pth', help='checkpoint to restore')
    parser.add_argument('--network', default='raft3d.raft3d', help='network architecture')
    args = parser.parse_args()

    torch_model = RAFT3D(args).to(device)

    CheckPoint = torch.load(model_path, weights_only=True, map_location=device)

    if 'state_dict' in CheckPoint:
        state_dict = CheckPoint['state_dict']
    else:
        state_dict = CheckPoint

    # Remove the 'module.' prefix if needed
    state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}
    missing_keys, unexpected_keys = torch_model.load_state_dict(state_dict, strict=False)
    
    # Set to eval mode
    torch_model.eval()
    
    fx, fy, cx, cy = (1050.0, 1050.0, 480.0, 270.0)
    image_previous = torch.randn(1, 3, 400, 640).cuda() # BatchSize = 1, Channel = 3, Height = 400, Width = 640
    image_current = torch.randn(1, 3, 400, 640).cuda() # BatchSize = 1, Channel = 3, Height = 400, Width = 640
    depth_previous = torch.randn(1, 400, 640).cuda() # BatchSize = 1, Channel = 1, Height = 400, Width = 640
    depth_current = torch.randn(1, 400, 640).cuda() # BatchSize = 1, Channel = 1, Height = 400, Width = 640
    intrinsics = torch.as_tensor([fx, fy, cx, cy]).cuda().unsqueeze(0)

    print("Initial Parameters are Finished")
    with torch.inference_mode():
        torch.onnx.export(
            torch_model, 
            (image_previous,image_current,depth_previous,depth_current,intrinsics), 
            '/home/h3c/LuoJiaxuan/RAFT3D/pretrained/raft3d.onnx',
            export_params=True,
            opset_version=16,
            do_constant_folding=True, 
            input_names = ["ImgPre","ImgCurr","DptPre","DptCurr","Intrinsic"],
            output_names = ["flow3d"],
            keep_initializers_as_inputs=False
            verbose=False,
        )
    print("Model Convert Onnx Finished")

Python Version: 3,9
Package Version:
cuda11.8
torch 2.0.1
torchaudio 2.0.2
torchvision 0.15.2
numpy 1.23.2
onnx 1.16.2
onnxruntime 1.18.1

I got a UserWarning: ONNX export mode is set to TrainingMode.EVAL, but operator 'instance_norm' is set to train=True. Exporting with train=True

Have encountered this issue and figured out a way to tackle this?
Thanks for your time.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

1 participant