Frame Averaging
Project description
Frame Averaging - Pytorch (wip)
Pytorch implementation of a simple way to enable (Stochastic) Frame Averaging for any network. This technique was recently adopted by Prescient Design in AbDiffuser
Install
$ pip install frame-averaging-pytorch
Usage
import torch
from frame_averaging_pytorch import FrameAverage
# contrived neural network
net = torch.nn.Linear(3, 3)
# wrap the network with FrameAverage
net = FrameAverage(
net,
dim = 3, # defaults to 3 for spatial, but can be any value
stochastic = True # whether to use stochastic variant from FAENet (one frame sampled at random)
)
# pass your input to the network as usual
points = torch.randn(2, 4, 1024, 3)
out = net(points)
out.shape # (2, 4, 1024, 3)
# frame averaging is automatically taken care of, as though the network were unwrapped
Citations
@article{Puny2021FrameAF,
title = {Frame Averaging for Invariant and Equivariant Network Design},
author = {Omri Puny and Matan Atzmon and Heli Ben-Hamu and Edward James Smith and Ishan Misra and Aditya Grover and Yaron Lipman},
journal = {ArXiv},
year = {2021},
volume = {abs/2110.03336},
url = {https://api.semanticscholar.org/CorpusID:238419638}
}
@article{Duval2023FAENetFA,
title = {FAENet: Frame Averaging Equivariant GNN for Materials Modeling},
author = {Alexandre Duval and Victor Schmidt and Alex Hernandez Garcia and Santiago Miret and Fragkiskos D. Malliaros and Yoshua Bengio and David Rolnick},
journal = {ArXiv},
year = {2023},
volume = {abs/2305.05577},
url = {https://api.semanticscholar.org/CorpusID:258564608}
}
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
frame_averaging_pytorch-0.0.10.tar.gz
(220.2 kB
view hashes)
Built Distribution
Close
Hashes for frame_averaging_pytorch-0.0.10.tar.gz
Algorithm | Hash digest | |
---|---|---|
SHA256 | 84a085c3a3d6f7abb10b9b5a8bad9167657c92f481f676580c14b3dd673a8d21 |
|
MD5 | 1ae3d5423cf450aeb46e42b0ce404eef |
|
BLAKE2b-256 | f5ab5fb2d6ab2dfa5b04477e22fbdab65be7949ce028449816de375e200e402c |
Close
Hashes for frame_averaging_pytorch-0.0.10-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 124e01bde22b68e3bca0133fa9b92e00e403c7b68509bef01bdd2b66fbf6b6ea |
|
MD5 | e4b6cdb92ec562141419e4d0524b2f6d |
|
BLAKE2b-256 | 5e2ada037ad6fb60482231fc56cd358160e68ab3f965cc3ca7c2087d4482fc8e |