|
|
--- |
|
|
title: README |
|
|
emoji: ⚡ |
|
|
colorFrom: blue |
|
|
colorTo: green |
|
|
sdk: static |
|
|
pinned: false |
|
|
--- |
|
|
|
|
|
# MLX KAN |
|
|
|
|
|
A community org for model weights compatible with `mlx-kan` powered by MLX. |
|
|
|
|
|
GitHub link: https://github.com/Goekdeniz-Guelmez/mlx-kan |
|
|
|
|
|
These are weights converted from pytorch and ready to be used. |
|
|
|
|
|
## How to install |
|
|
|
|
|
``` |
|
|
pip install mlx-kan |
|
|
``` |
|
|
|
|
|
## Models |
|
|
|
|
|
To load a model with pre-trained weights or create one from scratch: |
|
|
```python |
|
|
from mlx_kan.kan import KAN |
|
|
|
|
|
# Initialize and use KAN |
|
|
kan_model = KAN([in_features * out_features] + [hidden_dim] * (num_layers - 1) + [num_classes]) |
|
|
|
|
|
def train(model, train_set, train_labels, num_epochs=100): |
|
|
optimizer = optim.AdamW(learning_rate=0.0004, weight_decay=0.003) # Initialize a new optimizer for each model |
|
|
loss_and_grad_fn = nn.value_and_grad(model, loss_fn) |
|
|
|
|
|
# For 1 step |
|
|
loss, grads = loss_and_grad_fn(model, train_set, train_labels) |
|
|
optimizer.update(model, grads) |
|
|
mx.eval(model.parameters(), optimizer.state) |
|
|
avg_loss = total_loss += loss.item() |
|
|
|
|
|
# Update grid points here |
|
|
for name, layer in model.__dict__.items(): |
|
|
if isinstance(layer, KANLinear): |
|
|
with mx.no_grad(): |
|
|
layer.update_grid(train_set) |
|
|
``` |
