search

chenjianqu / RAFT_CPP

The optical flow algorithm RAFT implemented with C++(Libtorch+TensorRT)
MIT License
40 stars 4 forks source link

Single export (RAFT in full) possible? #4

Open christian1steinmetz opened 1 year ago

christian1steinmetz commented 1 year ago

Hi there,

I didn't find any other possibility to contact you, so I created this issue. Could you let me know why you decided to do three separate exports of the three submodels? Wouldn't it theoretically be possible to export the full RAFT model as it is and then let it run completely on a GPU?

I am asking as I am trying this at the moment, but am running in some problems during C++ Libtorch inference. Did you try this as well and did you run into similar problems?

Best Christian

itberrios commented 10 months ago

I haven't tried to do anything with this project yet, but maybe the 3 exports are due to the 3 main blocks of RAFT. RAFT is not just a straight feed forward neural net, it has the feature extraction block (f_net), the context encoder (c_net) and the update block.

Here's a (very) high level overview of how RAFT works:

Cai-RS commented 6 months ago

I haven't tried to do anything with this project yet, but maybe the 3 exports are due to the 3 main blocks of RAFT. RAFT is not just a straight feed forward neural net, it has the feature extraction block (f_net), the context encoder (c_net) and the update block.

Here's a (very) high level overview of how RAFT works:

  • f_net extracts features from frame 1 and frame 2
  • c_net extracts features from frame 1 (reference frame)
  • The Correlation Pyramid is obtained from the features from f_net
  • The Lookup operator samples features from the Correlation Pyramid
  • Both the c_net features and the sampled correlation features are passed to the update block for iterative refinement
  • And finally the output is upsampled Via Convex Upsampling.

The ovweview of RAFT is right. But why a network with multiple networks as components can not be exported as ONNX? RAFT itself is also a network, it must can be transfered once all its op are defined in torch.onnx opset