{"id":18676376,"url":"https://github.com/cgarciae/ciclo","last_synced_at":"2025-04-12T02:13:04.182Z","repository":{"id":61885686,"uuid":"548934741","full_name":"cgarciae/ciclo","owner":"cgarciae","description":"A functional training loops library for JAX","archived":false,"fork":false,"pushed_at":"2024-02-13T23:57:32.000Z","size":3669,"stargazers_count":86,"open_issues_count":10,"forks_count":8,"subscribers_count":3,"default_branch":"main","last_synced_at":"2025-04-12T02:12:56.410Z","etag":null,"topics":["jax","python"],"latest_commit_sha":null,"homepage":"https://cgarciae.github.io/ciclo","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/cgarciae.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2022-10-10T12:17:50.000Z","updated_at":"2024-12-10T20:28:00.000Z","dependencies_parsed_at":"2024-11-07T09:33:34.824Z","dependency_job_id":"f2163716-a4e9-4cb9-b436-fcde807fda25","html_url":"https://github.com/cgarciae/ciclo","commit_stats":null,"previous_names":[],"tags_count":2,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/cgarciae%2Fciclo","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/cgarciae%2Fciclo/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/cgarciae%2Fciclo/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/cgarciae%2Fciclo/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/cgarciae","download_url":"https://codeload.github.com/cgarciae/ciclo/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":248505925,"owners_count":21115354,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["jax","python"],"created_at":"2024-11-07T09:28:59.568Z","updated_at":"2025-04-12T02:13:04.150Z","avatar_url":"https://github.com/cgarciae.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"[![codecov](https://codecov.io/gh/cgarciae/ciclo/branch/main/graph/badge.svg?token=3IKEUAU3C8)](https://codecov.io/gh/cgarciae/ciclo)\n\n# 🌀 Ciclo\n_A functional training loops library for JAX_\n\n`ciclo` provides a set of utilities and abstractions to build complex training loops with any JAX framework. `ciclo` defines a set of building blocks that naturally compose together and scale up to build higher-level abstractions, ranging from low-level custom training loops to Keras-like training APIs.\n\n**Features**\n\n✔️ Training utilities \u003cbr\u003e\n✔️ Loop language \u003cbr\u003e\n✔️ Predefined Loops \u003cbr\u003e\n🧪 Managed API (simplified training + parallelism support) [**experimental**] \u003cbr\u003e\n🧪 Framework support (predifined states) [**experimental**] \u003cbr\u003e\n\n## Status\nCiclo is still in early development, the API is subject to change, expect things to break. If you are interested in contributing, please reach out.\n\n## Getting Started\nTo get started with Ciclo, you can install it via pip:\n\n```python\npip install ciclo\n```\nOnce you've installed Ciclo, you can import it into your Python script:\n\n```python\nimport ciclo\n```\n  \n### 🌀 Loop Language\n\nThe `loop` function in `ciclo` serves as a mini-language for defining training loops by composing functions. With the `tasks` dictionary, you can express the desired behavior of the loop as a composition of schedules and their corresponding callbacks.\n\nTo use the `loop` function, you first define your training steps as JAX functions, and then create a dictionary where the keys are schedules, and the values are lists of callbacks to execute at each scheduled interval.\n\n```python\n@jax.jit\ndef train_step(state, batch):\n    ... # do JAX stuff to update state\n    logs = ciclo.logs()\n    logs.add_metric(\"accuracy\", accuracy)\n    logs.add_metric(\"loss\", loss)\n    return logs, state\n\ntotal_steps = 100\nstate = create_state() # initial state\n\nstate, history, elapsed = ciclo.loop(\n    state, # Pytree\n    dataset, # Iterable[Batch]\n    { # Schedule: List[Callback]\n        ciclo.every(1): [train_step],\n        ciclo.every(steps=10): [ciclo.checkpoint(f\"logdir/model\")],\n        ciclo.every(1): [ciclo.keras_bar(total=total_steps)],\n    },\n    stop=total_steps,\n)\n```\n```\n80/100 [=====================\u003e.....] - ETA: 42s - accuracy: 0.6148 - loss: 1.537120\n```\n\nAt each iteration, callbacks can update the state and append new logs, the `loop` function returns the final state, the history of logs, and the elapsed time. Depending on the nature of each callback, the order in which they are executed may be very important e.g. `keras_bar` should always be last so that it can display the metrics produced by previous callbacks.\n\n### Predefined Loops\n\n`ciclo` provides a set of predefined loops that you can use out of the box for common scenarios:\n\n* `train_loop`: a training loop with an inner evaluation loop\n* `test_loop`: an evaluation loop\n* `predict_loop`: an inference loop\n\nAll of these loops are built on top of the `loop` function and extend the loop language with additional **named schedules** which run at specific points in the loop. They also provide a `callbacks` argument allows you to add callbacks without specifying a schedule, instead, if they contain an attribute that matches the name of a named schedule it will be automatically registered to that schedule. This also applies for methods of the `state` object.\n\nHere's an example of how to use the `train_loop`:\n\n```python\n@jax.jit\ndef test_step(state, batch):\n    ... # do JAX stuff\n    logs = ciclo.logs()\n    logs.add_metric(\"accuracy\", accuracy)\n    logs.add_metric(\"loss\", loss)\n    return logs, state\n\nstate, history, elapsed = ciclo.train_loop(\n    state, # Pytree\n    train_dataset, # Iterable[Batch]\n    { # Schedule: List[Callback]\n        ciclo.on_train_step: [train_step], # named schedules\n        ciclo.on_test_step: [test_step], # named schedules\n        ciclo.every(20): [some_callback], # regular schedules also supported\n    },\n    test_dataset=lambda: get_test_dataset(), # lazy test dataset definition\n    epoch_duration=steps_per_epoch,\n    callbacks=[\n        # callback self-registration\n        ciclo.keras_bar(total=total_steps), # runs on ciclo.on_train_batch_end\n        ciclo.checkpoint(f\"logdir/model\"),  # runs on ciclo.on_epoch_end\n    ],\n    stop=total_steps,\n)\n```\n```\n80/100 [=====================\u003e.....] - ETA: 42s - accuracy: 0.6148 - loss: 1.537120\n```\n\n### Managed API 🧪\nThe `managed` API aims to simplify the process of creating JAX programs for common patterns, such as `jit`, data parallelism with `pmap`, etc. To use this API, you need to define a compatible state type, which can be easily achieved by creating an instance of `managed.ManagedState` or a subclass of it.\n\n```python\nfrom ciclo import managed\n\nstate = managed.ManagedState.create(\n    params=params, # can be any pytree\n    tx=optax.adamw(1e-3), # optax optimizer\n    strategy=\"jit\", # \"data-parallel\" or \"eager\"\n)\n```\n\nWith the managed API, you can use the `train_step` decorator to define a training step easily. The `managed.train_step` function expects you to return logs with at least one loss, which will be used to automatically compute the gradients and update the parameters.\n\n```python\n@managed.train_step\ndef train_step(state, batch):\n    loss = ... # compute loss\n    logs = ciclo.logs()\n    logs.add_loss(\"loss\", loss) # \u003c\u003c\u003c register loss\n    return logs, state\n\nfor batch in train_dataset:\n    logs, state = train_step(state, batch)\n    print(f\"loss: {logs.losses.loss}\")\n```\n\nIf you need to perform some computation under a strategy without automatically computing gradients and updating the parameters, you can use the `managed.step` decorator.\n\n```python\n@managed.step\ndef test_step(state, batch):\n    loss = ... # compute loss\n    logs = ciclo.logs()\n    logs.add_metric(\"loss\", loss)\n    return logs, state\n```\n\n### Framework Support 🧪\n\nTo make JAX accessible for begginers, `ciclo` plans to provide a set of predefined state types for common frameworks, such as `flax`, `haiku`, and `equinox` (for now only `flax` is supported). These types are based on `ManagedState` and make it easy to perform tasks like training, evaluation, and inference without having to have a deep understanding of either the framework or JAX. Similar to Keras you just provide a model, an optimizer, some losses and metrics:\n\n```python\nimport flax.linen as nn\nfrom jax_metrics.losses import Crossentropy\nfrom jax_metrics.metrics import Accuracy\n\nmodel = nn.Sequential([\n    lambda x: x.reshape((x.shape[0], -1)) / 255.0,\n    nn.Dense(128),\n    nn.relu,\n    nn.Dense(10),\n])\nstate = ciclo.create_flax_state(\n    model,\n    inputs=jnp.empty((1, 28, 28, 1)),\n    tx=optax.adamw(1e-3),\n    losses={\"loss\": Crossentropy()}, # any Callable[..., jax.Array]\n    metrics={\"accuracy\": Accuracy()}, # supports Callables, jax_metrics, or clu metrics\n    strategy=\"jit\", # \"data-parallel\" or \"eager\"\n)\n```\n\nThese custom `state` objects provide `train_step`, `test_step`, and `predict_step` methods, this means that when used with the `train_loop`, `test_loop`, and `predict_loop` APIs, they provide a highly simplified experience:\n\n```python\nstate, history, elapsed = ciclo.train_loop(\n    state, # methods are automatically registered\n    train_dataset,\n    callbacks=[\n        ciclo.keras_bar(total=total_steps),\n    ],\n    test_dataset=lambda: get_test_dataset(),\n    epoch_duration=steps_per_epoch,\n    stop=total_steps,\n)\n```\n\nYou can also use the `train_step`, `test_step`, and `predict_step` methods directly to create custom training procedures:\n\n```python\nfor batch in train_dataset:\n    logs, state = state.train_step(batch)\n    print(f\"loss: {logs.losses.loss}\")\n```\n\n\n### Manual Iteration\nCiclo provides a set of loosely coupled APIs that can be used independently from the `loop` API to create custom training procedures when more control is needed. In the example below, we demonstrate how to manually iterate through a training dataset using `ciclo`.\n\nFirst, we define a few callbacks and utilities such as `call_checkpoint`, `checkpoint`, `keras_bar`, and `history`. Then, we use the `ciclo.elapse` function to iterate through the training dataset for a specified number of steps. During each iteration, we update the `logs` and `state` using the `train_step` function. We periodically checkpoint the `state`, update a progress bar, and commit the logs to the `history`.\n\n```python\ncall_checkpoint = ciclo.every(steps=1000) # Schedule\ncheckpoint = ciclo.checkpoint(f\"logdir/model\") # Callback\nkeras_bar = ciclo.keras_bar(total=total_steps) # Callback\nhistory = ciclo.history() # History\n# (Elapsed, Batch)\nfor elapsed, batch in ciclo.elapse(train_dataset, stop=total_steps):\n    logs = ciclo.logs() # Logs\n    # update logs and state\n    logs.updates, state = train_step(state, batch)\n    # periodically checkpoint state\n    if call_checkpoint(elapsed):\n        checkpoint(elapsed, state) # serialize state\n\n    keras_bar(elapsed, logs) # update progress bar\n    history.commit(elapsed, logs) # commit logs to history\n```\nThis approach allows for fine-grained control over the training process and enables customization of various aspects of the training loop.\n\n## Examples\n\nFor a more in-depth look at how to use `ciclo`, check out our [examples](./examples) folder which contains a set of python scripts that demonstrate how to use `ciclo` to train models using different APIs.\n\n### JAX\n* [00 Linear Regression](https://github.com/cgarciae/ciclo/blob/main/examples/jax/00_linear_regression_pure_jax.py)\n\n### Flax\n* [01 Simple MNIST](https://github.com/cgarciae/ciclo/blob/main/examples/flax/01_mnist_simple.py)\n* [02 Using train_loop](https://github.com/cgarciae/ciclo/blob/main/examples/flax/02_mnist_train_loop.py)\n* [03 Manual Iteration](https://github.com/cgarciae/ciclo/blob/main/examples/flax/03_mnist_manual_iteration.py)\n* [04 Managed API](https://github.com/cgarciae/ciclo/blob/main/examples/flax/04_mnist_managed_api.py)\n* [05 Using create_flax_state](https://github.com/cgarciae/ciclo/blob/main/examples/flax/05_mnist_flax_state.py)\n\n### Haiku\n* [Simple MNIST](https://github.com/cgarciae/ciclo/blob/main/examples/haiku/01_mnist_haiku.py)\n\n### Equinox\n* [Simple MNIST](https://github.com/cgarciae/ciclo/blob/main/examples/equinox/01_mnist_equinox.py)\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fcgarciae%2Fciclo","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fcgarciae%2Fciclo","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fcgarciae%2Fciclo/lists"}