Metadata-Version: 2.4
Name: iteryne
Version: 0.1.0
Summary: Model-Agnostic Meta-Learning (MAML) for any PyTorch nn.Module.
Project-URL: Homepage, https://github.com/wonderwice/iteryne
Project-URL: Documentation, https://github.com/wonderwice/iteryne#readme
Project-URL: Repository, https://github.com/wonderwice/iteryne
Author-email: Alexei Czornyj <alexei.czornyj@etu.univ-poitiers.fr>
License: MIT
License-File: LICENSE
Keywords: anil,few-shot,fomaml,maml,meta-learning,meta-sgd,pytorch
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Typing :: Typed
Requires-Python: >=3.10
Requires-Dist: torch>=2.1
Provides-Extra: dev
Requires-Dist: mypy>=1.10; extra == 'dev'
Requires-Dist: pre-commit>=3.5; extra == 'dev'
Requires-Dist: pytest-cov>=4.1; extra == 'dev'
Requires-Dist: pytest>=7.4; extra == 'dev'
Requires-Dist: ruff>=0.6; extra == 'dev'
Provides-Extra: docs
Requires-Dist: mkdocs-material>=9.5; extra == 'docs'
Requires-Dist: mkdocstrings[python]>=0.25; extra == 'docs'
Description-Content-Type: text/markdown

# Iteryne

Model-Agnostic Meta-Learning (MAML) for **any** PyTorch `nn.Module`.

`iteryne` is a small, well-tested implementation of MAML
([Finn, Abbeel & Levine, 2017](https://arxiv.org/abs/1703.03400)) built on
PyTorch's native `torch.func`. It meta-trains a model's initial parameters so
that a few gradient steps on a new task's small support set generalize to that
task's query set.

## Features

- **Model-agnostic.** Works on any `nn.Module` (MLPs, CNNs, BatchNorm, ...) with
  no rewriting, via `torch.func.functional_call`.
- **First- and second-order.** Full MAML (gradient-through-gradient) and FOMAML
  (the first-order approximation) share one code path; flip `first_order=True`.
- **Variants included.** Meta-SGD (learnable per-parameter inner learning rates)
  and ANIL (adapt only the head).
- **Two levels of API.** A high-level `MAML` wrapper plus `MetaTrainer`, and an
  exposed functional core (`adapt`, `inner_step`, `functional_forward`).
- **Typed and tested.** `py.typed`, mypy-strict, pytest suite including a
  `gradcheck` of the second-order meta-gradient.

## Install

```bash
pip install iteryne
# or, from a clone:
pip install -e ".[dev]"
```

Requires Python >= 3.10 and PyTorch >= 2.1.

## Quickstart

```python
import torch
from torch import nn
from iteryne import MAML, MetaTrainer, SinusoidTaskSampler

model = nn.Sequential(nn.Linear(1, 40), nn.ReLU(), nn.Linear(40, 40), nn.ReLU(), nn.Linear(40, 1))
maml = MAML(model, inner_lr=0.01, inner_steps=1, first_order=False)
meta_opt = torch.optim.Adam(maml.parameters(), lr=1e-3)

trainer = MetaTrainer(maml, meta_opt, nn.MSELoss(), SinusoidTaskSampler(seed=0))
trainer.fit(num_iterations=2000, meta_batch_size=25)

# Fast-adapt to a new task with a few gradient steps:
task = SinusoidTaskSampler(seed=123).sample(1)[0]
learner = maml.clone()
learner.adapt_on(nn.MSELoss(), task.support_x, task.support_y)
preds = learner(task.query_x)
```

### Manual training loop

`MetaTrainer` is optional. The core pattern:

```python
maml = MAML(model, inner_lr=0.01, inner_steps=5, first_order=False)
meta_opt = torch.optim.Adam(maml.parameters(), lr=1e-3)

meta_opt.zero_grad()
for task in task_batch:
    learner = maml.clone()
    for _ in range(maml.inner_steps):
        learner.adapt(loss_fn(learner(task.support_x), task.support_y))
    loss_fn(learner(task.query_x), task.query_y).backward()  # accumulates meta-grad
meta_opt.step()
```

The single `backward()` call computes the full second-order meta-gradient when
`first_order=False`, and the first-order (FOMAML) meta-gradient when `True`.

### Variants

```python
from iteryne import MetaSGD, ANIL

meta_sgd = MetaSGD(model, inner_lr=0.01)            # learns per-parameter inner LRs
anil = ANIL(model, head=model[-1], inner_lr=0.01)   # adapts only the head
```

## How it works

Parameters are carried as a `dict[str, Tensor]` and applied with
`functional_call`, so the module itself never has to be modified. The inner step

```
p' = p - alpha * grad(support_loss, p)
```

is computed with `torch.autograd.grad(..., create_graph=not first_order)`. With
`create_graph=True` the adapted parameters stay connected to the originals, so a
later `backward()` differentiates through the inner step (full MAML). With
`create_graph=False` the gradient term is detached, leaving the identity edge
`p' = p - alpha * g` so the same `backward()` yields the first-order meta-gradient
(FOMAML). One flag, one code path.

BatchNorm and other buffers are carried through the functional forward but are
not adapted in the inner loop; for few-shot work consider
`track_running_stats=False`.

## Documentation

Build the docs site locally:

```bash
pip install -e ".[docs]"
mkdocs serve
```

## Citation

```bibtex
@inproceedings{finn2017maml,
  title     = {Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks},
  author    = {Finn, Chelsea and Abbeel, Pieter and Levine, Sergey},
  booktitle = {International Conference on Machine Learning (ICML)},
  year      = {2017}
}
```
