sketchgraphs_models.torch_extensionsΒΆ

Pytorch extensions for our models.

This module provides some important extensions that work with segmented data, which is used extensively in our models due to the nature of graph batching. This modules provides a pure python implementation, although can leverage a C++ extension if present. This enables speed-up of over 2x (on GPU) when it is present, so it is recommended to compile the C++ extensions if possible.

Modules

sketchgraphs_models.torch_extensions.index

sketchgraphs_models.torch_extensions.segment_ops

sketchgraphs_models.torch_extensions.segment_pool