mindrally/machine-learning

Que fait machine-learning ?

Développement d'apprentissage automatique avec JAX, schémas de programmation fonctionnelle et calcul haute performance.

Aperçu

Ce skill guide les agents IA dans le développement de solutions d'apprentissage automatique efficaces en utilisant JAX, en mettant l'accent sur la programmation fonctionnelle et le calcul haute performance. Il aide les modèles à comprendre les structures de code existantes et à adhérer aux meilleures pratiques pour l'apprentissage profond, la différentiation automatique et la compilation JIT.

Utilisation

Lorsque vous travaillez sur des tâches connexes, le skill est activé automatiquement. Utilisez les directives suivantes :

1. Utiliser JAX : Privilégier JAX pour la différentiation automatique, la compilation JIT et l'accélération GPU/TPU.
2. Programmation fonctionnelle : Maintenir des fonctions pures, utiliser des transformations et éviter les effets de bord.
3. Calcul haute performance : Optimiser pour l'efficacité mémoire et la vitesse de calcul.

Directives d'implémentation

Gestion des états

• Utiliser flax.struct ou optax pour la gestion des états.
• Éviter les variables globales ; passer les paramètres explicitement.

Optimisation des gradients

import jax
import jax.numpy as jnp

def compute_gradients(params, x, y, model_fn):
    def loss_fn(p):
        return jnp.mean((model_fn(p, x) - y) ** 2)
    return jax.grad(loss_fn)(params)

Compilation JIT

import jax

def update_step(params, x, y, model_fn, learning_rate):
    grads = compute_gradients(params, x, y, model_fn)
    return params - learning_rate * grads

# JIT compile pour performance
update_step_jit = jax.jit(update_step)

Exemple d'architecture modulaire

import jax.numpy as jnp
import flax.linen as nn

class MLP(nn.Module):
    hidden_dims: list = [512, 256]
    num_classes: int = 10

    @nn.compact
    def __call__(self, x):
        x = x.reshape((x.shape[0], -1))  # aplatir
        for dim in self.hidden_dims:
            x = nn.Dense(dim)(x)
            x = nn.relu(x)
        x = nn.Dense(self.num_classes)(x)
        return x

Bonnes pratiques

• Utiliser jax.random.PRNGKey pour la génération de nombres aléatoires.
• Mettre en cache les opérations coûteuses avec jax.lax.map et jax.lax.scan.
• Préférer jax.grad pour les dérivées et jax.hessian pour les Hessiennes.
• Assurer la reproductibilité en définissant des clés aléatoires et des lots déterministes.

Dépannage

• Si des problèmes de mémoire surviennent, envisager d'utiliser jax.lax.reduce ou jax.lax.fori_loop pour l'évaluation par lots.
• Pour des séquences longues, utiliser des transformations jax.lax.scan au lieu de boucles Python.

Références

• Documentation JAX
• Documentation Flax
• Documentation Optax