Skip to content

Implementations of the custom layers in C++ and Native Pytorch for CPU support.#212

Open
zsameem wants to merge 2 commits into
NVIDIA:masterfrom
zsameem:cpu_support
Open

Implementations of the custom layers in C++ and Native Pytorch for CPU support.#212
zsameem wants to merge 2 commits into
NVIDIA:masterfrom
zsameem:cpu_support

Conversation

@zsameem

@zsameem zsameem commented May 16, 2020

Copy link
Copy Markdown
Contributor

Implementation of Correlation, Resample2D layer and the ChannelNorm layer in native Pytorch and C++ to support inference on CPU.

The main bottleneck is the Correlation Layer on which the FlowNetC architecture relies.
This PR provides 2 implementations of the Correlation layer.

-PyTorch native implementation. This requires no extra setup
-Optimized C++ implementation for inference on CPU.

Also provided are Pytorch native implementations for Resample2D and Channelnorm.
Since the Pytorch implementation is quite efficient (compeletely vectorized) with no
python for loops, C++ implementation is not needed. These layers also work by default
on the GPU dependening on if the input tensors are on gpu and are slightly slower than
the provided cuda implementation.

See comments at the top of models.py and networks/FlowNetC.py for more details and
how to switch to CPU mode.

Backward passes are not yet implemented but will be added in the future.

run_a_pair.py is replaced with a generic script called test.py to simply test functionality.
run_a_pair.py had hardcoded paths. Also 2 frames from sintel added in test_images dir so that
functionality and setup can be swiftly checked.

Resolves: #190

zsameem added 2 commits May 16, 2020 17:04
…ple2D layer and the ChannelNorm layer in native Pytorch and C++ to support inference on CPU.

The main bottleneck is the Correlation Layer on which the FlowNetC architecture relies.
This PR provides 2 implementations of the Correlation layer.

-PyTorch native implementation. This requires no extra setup
-Optimized C++ implementation for inference on CPU.

Also provided are Pytorch native implementations for Resample2D and Channelnorm.
Since the Pytorch implementation is quite efficient (compeletely vectorized) with no
python for loops, C++ implementation is not needed. These layers also work by default
on the GPU dependening on if the input tensors are on gpu and are slightly slower than
the provided cuda implementation.

See comments at the top of models.py and networks/FlowNetC.py for more details and
how to switch to CPU mode.

Backward passes are not yet implemented but will be added in the future.

run_a_pair.py is replaced with a generic script called test.py to simply test functionality.
run_a_pair.py had hardcoded paths. Also 2 frames from sintel added in test_images dir so that
functionality and setup can be swiftly checked.

Resolves: NVIDIA#190
@Ahleroy

Ahleroy commented Nov 26, 2020

Copy link
Copy Markdown

I do not see backward method for the Correlation implementations (PyTorch and C++).

EDIT : I mean not completed

@zsameem

zsameem commented Nov 26, 2020

Copy link
Copy Markdown
Contributor Author

I do not see backward method for the Correlation implementations (PyTorch and C++).

EDIT : I mean not completed

I did not implement them because I only needed forward methods for inference.

@Ahleroy

Ahleroy commented Nov 27, 2020

Copy link
Copy Markdown

So they do not work for training ?
I see the PyTorch correlation is not compatible with the last version of PyTorch. Foward and backward functions must be static method.

@zsameem

zsameem commented Nov 28, 2020

Copy link
Copy Markdown
Contributor Author

Yes, they will not work for training. And yes, newer PyTorch versions introduced some API changes such as making forward and backward static and slightly different calling syntax for nn.Function classes. These have not been updated yet.

@batsheva-knecht

Copy link
Copy Markdown

Hi, I tried to export the model FlowNet2 to onnx (the one with Correlation, Resample2d and ChannelNorm in PyTorch version), However I get: RuntimeError: ONNX export failed: Couldn't export Python operator CorrelationFunction. Does someone know how to solve it?

@andreyiva111

Copy link
Copy Markdown

Hi!
@batsheva-knecht , did you manage to convert to the onnx format?
I have the same problem. (with one difference that it is a model PWCNet) The only thing I can think of is to convert the model piece by piece to the onnx format. And between calls to these parts apply correlation kernels cuda... Maybe there are other options?

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

Successfully merging this pull request may close these issues.

Resample2D function to CPU compatible code

4 participants