File size: 1,194 Bytes
39de75b 315b458 39de75b 315b458 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
---
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)
``` |