DOCSEARCHbyAlgolia
Submit Docs

escnn (E(n)-equivariant Steerable CNNs)

repository·master·Indexed 19 days ago

https://github.com/quva-lab/escnn

A PyTorch extension for building neural networks equivariant to symmetry groups, such as rotations and reflections, for isometries E(n) of R^n where n=2 and n=3. The library provides high-level abstractions for group and representation theory, utilizing GeometricTensors and feature fields to ensure geometrically sound processing. It includes subpackages for group theory (escnn.group), equivariant convolution kernels (escnn.kernels), Euclidean spaces (escnn.gspaces), and equivariant neural network modules (escnn.nn).

Tokens
12.6K
Snippets
25
Records
64
Agent score
68%

What's inside escnn

  1. Overview of escnn

    master

    escnn is a PyTorch-based library designed for equivariant deep learning. It provides tools for the equivariant processing of signals over Euclidean spaces, specifically supporting 2D and 3D signals (e.g., planar images or 3D volumes).

    Key capabilities include:

    • Implementing the most general convolutional maps equivariant under isometries (translations, rotations, and reflections) of Euclidean space.
    • Supporting compact-group equivariant linear maps (equivariant MLPs) by treating them as a special case of equivariant maps on a 0-dimensional Euclidean space.

  2. Overview of escnn subpackages

    master

    The library is organized into four main subpackages:

    ComponentDescription
    escnn.groupImplements basic concepts of group and representation theory
    escnn.kernelsSolves for spaces of equivariant convolution kernels
    escnn.gspacesDefines the Euclidean spaces and their symmetries
    escnn.nnContains equivariant modules to build deep neural networks

  3. What is escnn and how does it work?

    master

    escnn is a PyTorch extension for equivariant deep learning. It provides neural network modules that are equivariant under all isometries $\mathrm{E}(n)$ of $\mathbb{R}^n$ (translations, rotations, and reflections) for $n=2$ and $n=3$.

    Unlike conventional CNNs, $\mathrm{E}(n)$-equivariant models are guaranteed to generalize over these transformations, making them more data-efficient.

    Key Concepts:

    • Feature Fields: Instead of standard channels, feature spaces are defined as feature fields characterized by their transformation laws (e.g., scalar fields like grayscale images or vector fields like optical flow).
    • GeometricTensor: The library represents feature fields using GeometricTensor objects, which wrap a torch.Tensor with its corresponding transformation law. All equivariant operations perform dynamic type-checking to ensure geometrically sound processing.
    • Defining Feature Spaces: Users specify field types and their multiplicities rather than just a number of channels. Modules like R2conv and R3conv then instantiate the most general convolutional mapping between specified input and output feature spaces.

  4. Understand the escnn package structure

    master

    The library is organized into four main subpackages. While all are part of the ecosystem, most users building equivariant models will primarily interact with escnn.gspaces and escnn.nn.

    • escnn.group: Implements fundamental concepts of group and representation theory.
    • escnn.kernels: Solves for spaces of equivariant convolution kernels.
    • escnn.gspaces: Defines Euclidean spaces and their associated symmetries.
    • escnn.nn: Contains equivariant modules used to construct deep neural networks.

  5. Steerable Point Convolutions

    master

    These modules provide convolution capabilities for point clouds or non-grid data using a basis expansion approach.

    Key components include:

    • Convolution Modules: R2PointConv, R3PointConv, RdPointConv.
    • BasisManager: Manages the basis expansion used by point convolutions, including BlocksBasisExpansion and BlocksBasisSampler.

  6. Coordinate conventions for points and voxel data

    master

    Understanding the coordinate convention is critical when transforming GeometricTensor objects or building filters in R3Conv.

    1. Standard Point Convention: A point $\mathbf{v} \in \mathbb{R}^n$ is parameterized as $(X, Y, Z, \dots)^T$.
    2. Voxel Data Convention: When working with voxel data, escnn uses a $(\dots, -Z, -Y, X)$ convention. In a feature tensor of shape $(B, C, D_1, D_2, \dots, D_{n-2}, D_{n-1}, D_n)$, the last dimension is the $X$ axis, but the $i$-th last dimension is the inverted $i$-th axis. This is consistent with the $(-Y, X)$ convention used in 2D images.

  7. Manipulate representations in escnn.group

    master

    The escnn.group subpackage provides tools to combine and transform group representations.

    Key Operations

    • Direct Sum: Combine representations using escnn.group.direct_sum() or the binary operator +.
    • Tensor Product: Combine representations using the Representation.tensor() method.
    • Restriction: Restrict a representation to a subgroup using Representation.restrict().
    • Induction: Induce a representation to a larger group using Group.induced_representation().

    Core Classes

    • Representation: The base class for group representations.
    • IrreducibleRepresentation: Represents an irreducible representation of the group.

  8. Configure the steerable basis in R2Conv and R3Conv

    master

    The steerable basis for escnn.nn.modules.conv._RdConv (and its subclasses R2Conv or R3Conv) is controlled by three optional parameters. While the default values are tuned for odd-sized, small kernel filters and are recommended for most use cases, you can customize them for even-sized or very wide filters.

    • rings: A list of float values defining the radii of the different rings the filter is split into.
    • sigma: Defines the width of each ring. You can provide a single float (used for all rings) or a list of floats (one per radius).
    • frequencies_cutoff: Regulates the maximum frequency on each ring.
      • If a float F is provided, the maximum frequency at radius r is int(r * F).
      • Alternatively, you can pass a function that accepts a radius r and returns the maximum frequency for that radius.

  9. How to use the escnn.kernels subpackage

    master

    The escnn.kernels subpackage implements analytical solutions for equivariance constraints on kernel spaces.

    Important Usage Note: Typically, users should not interact with this subpackage directly. Instead, use the higher-level interface provided in escnn.gspaces. The kernels are intended to be built through factory functions.

    PyTorch Integration: Every KernelBasis is a subclass of torch.nn.Module. This means you can:

    • Move it to a device (e.g., .to('cuda')).
    • Change floating point precision.
    • Use it in a differentiable forward pass.

    Version Warning: This module was refactored in version 1.0.0. Weights trained with versions prior to 1.0.0 may be incompatible with newer model instantiations due to changes in basis element indexing.

  10. Use factory functions to instantiate groups in escnn.group

    master

    To avoid creating multiple redundant instances of the same group, use the factory functions provided in escnn.group. These functions implement a singleton-like pattern where they build a single instance of each different group and return that same instance on consecutive calls.

    While you can instantiate objects from this subpackage directly, it is generally recommended to use the higher-level interface provided in escnn.gspaces when building equivariant networks.

  11. How FieldType, GeometricTensor, and EquivariantModule work together

    master

    The escnn.nn subpackage implements equivariant neural network modules using three core abstractions:

    1. FieldType: Acts as a data type that defines the transformation law of a feature field under a symmetry group. It specifies how features behave when the input is transformed.
    2. GeometricTensor: A wrapper around a standard torch.Tensor that endows it with a FieldType. This allows the tensor to carry information about its geometric properties.
    3. EquivariantModule: A subclass of torch.nn.Module that processes GeometricTensors. It guarantees that the output field's transformation law is consistent with the input field's transformation law given a symmetry transformation.

    To build a model, you define FieldTypes for your inputs, hidden layers, and outputs, then compose EquivariantModules (like R2Conv or ReLU) into a network.

    s = escnn.gspaces.rot2dOnR2(8)
c_in = escnn.nn.FieldType(s, [s.trivial_repr]*3)
c_hid = escnn.nn.FieldType(s, [s.regular_repr]*3)
c_out = escnn.nn.FieldType(s, [s.regular_repr]*1)

# net is an EquivariantModule
net = SequentialModule(
    R2Conv(c_in, c_hid, 5, bias=False),
    InnerBatchNorm(c_hid),
    ReLU(c_hid, inplace=True),
    # ...
)

  12. How to define symmetries using G-spaces

    master

    In escnn, symmetries are defined using GSpace objects. A GSpace specifies the symmetries considered for a given space (like a plane or 3D volume). You typically use these GSpace instances to instantiate equivariant neural network modules in escnn.nn.

    While you can subclass escnn.gspaces.GSpace to define custom symmetries, it is recommended to use the provided factory functions for common symmetries. These functions offer a simpler interface for instantiating the most frequently used group actions.