bartz Télécharger - bartz Téléchargement du code source

bartz

Autre code source

v chain samples we made along the way

Télécharger

BART vectorisé

Une implémentation d'arbres de régression bayésiens additifs (BART) dans JAX.

Si vous ne savez pas ce qu'est BART, mais connaissez XGBoost, considérez BART comme une sorte de XGBoost bayésien. bartz fait fonctionner BART aussi vite que XGBoost.

BART est une technique de régression bayésienne non paramétrique. Compte tenu des prédicteurs de formation $X$ et les réponses $y$ , BART trouve une fonction pour prédire $y$ donné $X$ . Le résultat de l'inférence est un échantillon de fonctions possibles, représentant l'incertitude sur la détermination de la fonction.

Ce module Python fournit une implémentation de BART qui s'exécute sur GPU, pour traiter plus rapidement de grands ensembles de données. C'est également bon sur le CPU. La plupart des autres implémentations de BART sont destinées à R et fonctionnent uniquement sur CPU.

Sur CPU, bartz fonctionne à la vitesse de dbarts (l'implémentation la plus rapide que je connaisse) si n > 20 000, mais en utilisant 1/20 de la mémoire. Sur GPU, la prime de vitesse dépend de la taille de l'échantillon ; c'est pratique sur CPU uniquement pour n > 10 000. L'accélération maximale est actuellement de 200x, sur un Nvidia A100 et avec au moins 2 000 000 d'observations.

Ce notebook Colab exécute bartz avec n = 100 000 observations, p = 1 000 prédicteurs, 10 000 arbres, pour 1 000 itérations MCMC, en 5 minutes.