NVIDIA/torch-harmonics

解決する課題

torch-harmonics は、球面上の微分可能な信号処理のためのフレームワークを提供し、球形データを機械学習モデルに統合することを可能にします。これは、微分可能性や効率を損なうことなく球面上でフーリエ変換のような変換や畳み込みを実行するという困難な問題を解決します。これは、偏微分方程式 (PDE) の求解や球面ドメインでのニューラル演算子の学習において極めて重要です。

仕組み

このライブラリは、フーリエ変換を球面に一般化した球面調和関数変換 (SHT) を実装しています。求積法を用いて信号を随伴ルジャンドル多項式に投影し、高速フーリエ変換 (FFT) を用いて調和基底を計算します。ディープラーニングとの互換性を確保するため、PyTorchのプリミティブを使用して構築されており、すべての演算が完全に微分可能です。また、測地距離を使用して近傍を決定する局所離散連続 (DISCO) 畳み込みや球面アテンション機構を導入しており、演算が近似的に等変であることを保証します。

対象ユーザー

微分可能な PDE ソルバー、球面フーリエニューラル演算子 (SFNO)、または球面上の深度推定やセマンティックセグメンテーションモデルを開発する、球形データを扱う研究者やエンジニア。

ハイライト

• 微分可能な SHT: 球面、ベクトル球面、および離散連続畳み込みの完全に微分可能な実装。
• GPU 加速: NVIDIA GPU での高性能計算のために最適化された CUDA 拡張。
• 分散コンピューティング: 空間的に分散された求積をサポートし、スケーラブルなモデル並列化を実現。
• 球面アテンション: 測地距離を使用して近傍を決定し、アテンション機構を球面に一般化。
• 幅広い応用: 浅水方程式の求解、深度推定、球面グリッド上のセマンティックセグメンテーションなどのタスクに使用可能。