We read every piece of feedback, and take your input very seriously.
To see all available qualifiers, see our documentation.
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
Thanks to the author for your wonderful work. I'm going to deploy this model as a part of our project.
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
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
Have encountered this issue and figured out a way to tackle this? Thanks for your time.
The text was updated successfully, but these errors were encountered:
No branches or pull requests
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).
And my Convert2OnnxModel Function is followed as:
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
The text was updated successfully, but these errors were encountered: