| Name | mlfab JSON |
| Version |
0.2.9
JSON |
| download |
| home_page | https://github.com/dpshai/mlfab |
| Summary | A collection of core machine learning tools |
| upload_time | 2024-11-06 23:00:13 |
| maintainer | None |
| docs_url | None |
| author | Benjamin Bolte |
| requires_python | >=3.11 |
| license | None |
| keywords |
|
| VCS |
 |
| bugtrack_url |
|
| requirements |
No requirements were recorded.
|
| Travis-CI |
No Travis.
|
| coveralls test coverage |
No coveralls.
|
<p align="center">
<picture>
<img alt="K-Scale Open Source Robotics" src="https://media.kscale.dev/kscale-open-source-header.png" style="max-width: 100%;">
</picture>
</p>
<div align="center">
[](https://github.com/dpshai/mlfab/blob/master/LICENSE)
[](https://discord.gg/k5mSvCkYQh)
[](https://humanoids.wiki)
<br />
[](https://github.com/pre-commit/pre-commit)
[](https://pytorch.org/get-started/locally/)
[](https://black.readthedocs.io/en/stable/)
[](https://github.com/charliermarsh/ruff)
<br />
[](https://github.com/kscalelabs/mlfab/actions/workflows/test.yml)
</div>
<br />
# mlfab
## What is this?
This is a framework for trying out machine learning ideas.
## Getting Started
Install the package using:
```bash
pip install mlfab
```
Or, to install the latest branch:
```bash
pip install 'mlfab @ git+https://github.com/kscalelabs/mlfab.git@master'
```
### Simple Example
This framework provides an abstraction for quickly implementing and training PyTorch models. The workhorse for doing this is `mlfab.Task`, which wraps all of the training logic into a single cohesive unit. We can override functions on that method to get special functionality, but the default functionality is often good enough. Here's an example for training an MNIST model:
```python
from dataclasses import dataclass
import torch.nn.functional as F
from dpshdl.dataset import Dataset
from dpshdl.impl.mnist import MNIST
from torch import Tensor, nn
from torch.optim.optimizer import Optimizer
import mlfab
@dataclass(kw_only=True)
class Config(mlfab.Config):
in_dim: int = mlfab.field(1, help="Number of input dimensions")
learning_rate: float = mlfab.field(1e-3, help="Learning rate to use for optimizer")
betas: tuple[float, float] = mlfab.field((0.9, 0.999), help="Beta values for Adam optimizer")
weight_decay: float = mlfab.field(1e-4, help="Weight decay to use for the optimizer")
warmup_steps: int = mlfab.field(100, help="Number of warmup steps to use for the optimizer")
class MnistClassification(mlfab.Task[Config]):
def __init__(self, config: Config) -> None:
super().__init__(config)
self.model = nn.Sequential(
nn.Conv2d(config.in_dim, 32, 3, padding=1),
nn.BatchNorm2d(32),
nn.ReLU(),
nn.Conv2d(32, 32, 3, padding=1),
nn.BatchNorm2d(32),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 3, padding=1),
nn.BatchNorm2d(64),
nn.ReLU(),
nn.Conv2d(64, 64, 3, padding=1),
nn.BatchNorm2d(64),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Flatten(),
nn.Linear(64 * 7 * 7, 128),
nn.BatchNorm1d(128),
nn.ReLU(),
nn.Linear(128, 10),
)
def set_loggers(self) -> None:
self.add_logger(
mlfab.StdoutLogger(),
mlfab.TensorboardLogger(self.exp_dir),
)
def get_dataset(self, phase: mlfab.Phase) -> MNIST:
root_dir = mlfab.get_data_dir() / "mnist"
return MNIST(root_dir=root_dir, train=phase == "train")
def forward(self, x: Tensor) -> Tensor:
return self.model(x)
def get_loss(self, batch: tuple[Tensor, Tensor], state: mlfab.State) -> Tensor:
x_bhw, y_b = batch
x_bchw = (x_bhw.float() / 255.0).unsqueeze(1)
yhat_bc = self(x_bchw)
self.log_step(batch, yhat_bc, state)
return F.cross_entropy(yhat_bc, y_b)
def log_valid_step(self, batch: tuple[Tensor, Tensor], output: Tensor, state: mlfab.State) -> None:
(x_bhw, y_b), yhat_bc = batch, output
def get_label_strings() -> list[str]:
ypred_b = yhat_bc.argmax(-1)
return [f"ytrue={y_b[i]}, ypred={ypred_b[i]}" for i in range(len(y_b))]
self.log_labeled_images("images", lambda: (x_bhw, get_label_strings()))
if __name__ == "__main__":
# python -m examples.mnist
MnistClassification.launch(Config(batch_size=16))
```
Let's break down each part individually.
### Config
Tasks are parametrized using a config dataclass. The `ml.field` function is a lightweight wrapper around `dataclasses.field` which is a bit more ergonomic, and `ml.Config` is a bigger dataclass which contains a bunch of other options for configuring training.
```python
@dataclass(kw_only=True)
class Config(mlfab.Config):
in_dim: int = mlfab.field(1, help="Number of input dimensions")
```
### Model
All tasks should subclass `ml.Task` and override the generic `Config` with the task-specific config. This is very important, not just because it makes your life easier by working nicely with your typechecker, but because the framework looks at the generic type when resolving the config for the given task.
```python
class MnistClassification(mlfab.Task[Config]):
def __init__(self, config: Config) -> None:
super().__init__(config)
self.model = nn.Sequential(
nn.Conv2d(config.in_dim, 32, 3, padding=1),
nn.BatchNorm2d(32),
nn.ReLU(),
nn.Conv2d(32, 32, 3, padding=1),
nn.BatchNorm2d(32),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 3, padding=1),
nn.BatchNorm2d(64),
nn.ReLU(),
nn.Conv2d(64, 64, 3, padding=1),
nn.BatchNorm2d(64),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Flatten(),
nn.Linear(64 * 7 * 7, 128),
nn.BatchNorm1d(128),
nn.ReLU(),
nn.Linear(128, 10),
)
```
### Loggers
`mlfab` supports logging to multiple downstream loggers, and provides a bunch of helper functions for doing common logging operations, like rate limiting, converting image resolution to normal sizes, overlaying captions on images, and more.
If this function is not overridden, the task will just log to `stdout`.
```python
def set_loggers(self) -> None:
self.add_logger(
mlfab.StdoutLogger(),
mlfab.TensorboardLogger(self.exp_dir),
)
```
### Datasets
The task should return the dataset used for training, based on the phase. `ml.Phase` is a string literal with values in `["train", "valid"]`. `mlfab.get_data_dir()` returns the data directory, which can be set in a configuration file which lives in `~/.mlfab.yml`. The default configuration file will be written on first run if it doesn't exist yet.
```python
def get_dataset(self, phase: mlfab.Phase) -> Dataset[tuple[Tensor, Tensor]]:
root_dir = mlfab.get_data_dir() / "mnist"
return MNIST(root_dir=root_dir, train=phase == "train")
```
### Compute Loss
Each `mlfab` model should either implement the `forward` function, which should take a batch from the dataset and return the loss, or, if more control is desired, the `get_loss` function can be overridden.
```python
def forward(self, x: Tensor) -> Tensor:
return self.model(x)
def get_loss(self, batch: tuple[Tensor, Tensor], state: mlfab.State) -> Tensor:
x_bhw, y_b = batch
x_bchw = (x_bhw.float() / 255.0).unsqueeze(1)
yhat_bc = self(x_bchw)
self.log_step(batch, yhat_bc, state)
return F.cross_entropy(yhat_bc, y_b)
```
### Logging
When we call `log_step` in the `get_loss` function, it delegates to either `log_train_step` or `log_valid_step`, depending on what `state.phase` is. In this case, on each validation step we log images of the MNIST digits with the labels that our model predicts.
```python
def log_valid_step(self, batch: tuple[Tensor, Tensor], output: Tensor, state: mlfab.State) -> None:
(x_bhw, y_b), yhat_bc = batch, output
def get_label_strings() -> list[str]:
ypred_b = yhat_bc.argmax(-1)
return [f"ytrue={y_b[i]}, ypred={ypred_b[i]}" for i in range(len(y_b))]
self.log_labeled_images("images", lambda: (x_bhw, get_label_strings()))
```
### Running
We can launch a training job using the `launch` class method. The config can be a `Config` object, or it can be the path to a `config.yaml` file located in the same directory as the task file. You can additionally provide the `launcher` argument, which supports training the model across multiple GPUs or nodes.
```python
if __name__ == "__main__":
MnistClassification.launch(Config(batch_size=16))
```
Raw data
{
"_id": null,
"home_page": "https://github.com/dpshai/mlfab",
"name": "mlfab",
"maintainer": null,
"docs_url": null,
"requires_python": ">=3.11",
"maintainer_email": null,
"keywords": null,
"author": "Benjamin Bolte",
"author_email": null,
"download_url": "https://files.pythonhosted.org/packages/6b/a3/c7ee9791301b0b5232fdc2c6f6929d9550dc901dcfa0b25f7b59b468fe88/mlfab-0.2.9.tar.gz",
"platform": null,
"description": "<p align=\"center\">\n <picture>\n <img alt=\"K-Scale Open Source Robotics\" src=\"https://media.kscale.dev/kscale-open-source-header.png\" style=\"max-width: 100%;\">\n </picture>\n</p>\n\n<div align=\"center\">\n\n[](https://github.com/dpshai/mlfab/blob/master/LICENSE)\n[](https://discord.gg/k5mSvCkYQh)\n[](https://humanoids.wiki)\n<br />\n[](https://github.com/pre-commit/pre-commit)\n[](https://pytorch.org/get-started/locally/)\n[](https://black.readthedocs.io/en/stable/)\n[](https://github.com/charliermarsh/ruff)\n<br />\n[](https://github.com/kscalelabs/mlfab/actions/workflows/test.yml)\n\n</div>\n\n<br />\n\n# mlfab\n\n## What is this?\n\nThis is a framework for trying out machine learning ideas.\n\n## Getting Started\n\nInstall the package using:\n\n```bash\npip install mlfab\n```\n\nOr, to install the latest branch:\n\n```bash\npip install 'mlfab @ git+https://github.com/kscalelabs/mlfab.git@master'\n```\n\n### Simple Example\n\nThis framework provides an abstraction for quickly implementing and training PyTorch models. The workhorse for doing this is `mlfab.Task`, which wraps all of the training logic into a single cohesive unit. We can override functions on that method to get special functionality, but the default functionality is often good enough. Here's an example for training an MNIST model:\n\n```python\n\nfrom dataclasses import dataclass\n\nimport torch.nn.functional as F\nfrom dpshdl.dataset import Dataset\nfrom dpshdl.impl.mnist import MNIST\nfrom torch import Tensor, nn\nfrom torch.optim.optimizer import Optimizer\n\nimport mlfab\n\n\n@dataclass(kw_only=True)\nclass Config(mlfab.Config):\n in_dim: int = mlfab.field(1, help=\"Number of input dimensions\")\n learning_rate: float = mlfab.field(1e-3, help=\"Learning rate to use for optimizer\")\n betas: tuple[float, float] = mlfab.field((0.9, 0.999), help=\"Beta values for Adam optimizer\")\n weight_decay: float = mlfab.field(1e-4, help=\"Weight decay to use for the optimizer\")\n warmup_steps: int = mlfab.field(100, help=\"Number of warmup steps to use for the optimizer\")\n\n\nclass MnistClassification(mlfab.Task[Config]):\n def __init__(self, config: Config) -> None:\n super().__init__(config)\n\n self.model = nn.Sequential(\n nn.Conv2d(config.in_dim, 32, 3, padding=1),\n nn.BatchNorm2d(32),\n nn.ReLU(),\n nn.Conv2d(32, 32, 3, padding=1),\n nn.BatchNorm2d(32),\n nn.ReLU(),\n nn.MaxPool2d(2),\n nn.Conv2d(32, 64, 3, padding=1),\n nn.BatchNorm2d(64),\n nn.ReLU(),\n nn.Conv2d(64, 64, 3, padding=1),\n nn.BatchNorm2d(64),\n nn.ReLU(),\n nn.MaxPool2d(2),\n nn.Flatten(),\n nn.Linear(64 * 7 * 7, 128),\n nn.BatchNorm1d(128),\n nn.ReLU(),\n nn.Linear(128, 10),\n )\n\n def set_loggers(self) -> None:\n self.add_logger(\n mlfab.StdoutLogger(),\n mlfab.TensorboardLogger(self.exp_dir),\n )\n\n def get_dataset(self, phase: mlfab.Phase) -> MNIST:\n root_dir = mlfab.get_data_dir() / \"mnist\"\n return MNIST(root_dir=root_dir, train=phase == \"train\")\n\n def forward(self, x: Tensor) -> Tensor:\n return self.model(x)\n\n def get_loss(self, batch: tuple[Tensor, Tensor], state: mlfab.State) -> Tensor:\n x_bhw, y_b = batch\n x_bchw = (x_bhw.float() / 255.0).unsqueeze(1)\n yhat_bc = self(x_bchw)\n self.log_step(batch, yhat_bc, state)\n return F.cross_entropy(yhat_bc, y_b)\n\n def log_valid_step(self, batch: tuple[Tensor, Tensor], output: Tensor, state: mlfab.State) -> None:\n (x_bhw, y_b), yhat_bc = batch, output\n\n def get_label_strings() -> list[str]:\n ypred_b = yhat_bc.argmax(-1)\n return [f\"ytrue={y_b[i]}, ypred={ypred_b[i]}\" for i in range(len(y_b))]\n\n self.log_labeled_images(\"images\", lambda: (x_bhw, get_label_strings()))\n\n\nif __name__ == \"__main__\":\n # python -m examples.mnist\n MnistClassification.launch(Config(batch_size=16))\n```\n\nLet's break down each part individually.\n\n### Config\n\nTasks are parametrized using a config dataclass. The `ml.field` function is a lightweight wrapper around `dataclasses.field` which is a bit more ergonomic, and `ml.Config` is a bigger dataclass which contains a bunch of other options for configuring training.\n\n```python\n@dataclass(kw_only=True)\nclass Config(mlfab.Config):\n in_dim: int = mlfab.field(1, help=\"Number of input dimensions\")\n```\n\n### Model\n\nAll tasks should subclass `ml.Task` and override the generic `Config` with the task-specific config. This is very important, not just because it makes your life easier by working nicely with your typechecker, but because the framework looks at the generic type when resolving the config for the given task.\n\n```python\nclass MnistClassification(mlfab.Task[Config]):\n def __init__(self, config: Config) -> None:\n super().__init__(config)\n\n self.model = nn.Sequential(\n nn.Conv2d(config.in_dim, 32, 3, padding=1),\n nn.BatchNorm2d(32),\n nn.ReLU(),\n nn.Conv2d(32, 32, 3, padding=1),\n nn.BatchNorm2d(32),\n nn.ReLU(),\n nn.MaxPool2d(2),\n nn.Conv2d(32, 64, 3, padding=1),\n nn.BatchNorm2d(64),\n nn.ReLU(),\n nn.Conv2d(64, 64, 3, padding=1),\n nn.BatchNorm2d(64),\n nn.ReLU(),\n nn.MaxPool2d(2),\n nn.Flatten(),\n nn.Linear(64 * 7 * 7, 128),\n nn.BatchNorm1d(128),\n nn.ReLU(),\n nn.Linear(128, 10),\n )\n```\n\n### Loggers\n\n`mlfab` supports logging to multiple downstream loggers, and provides a bunch of helper functions for doing common logging operations, like rate limiting, converting image resolution to normal sizes, overlaying captions on images, and more.\n\nIf this function is not overridden, the task will just log to `stdout`.\n\n```python\ndef set_loggers(self) -> None:\n self.add_logger(\n mlfab.StdoutLogger(),\n mlfab.TensorboardLogger(self.exp_dir),\n )\n```\n\n### Datasets\n\nThe task should return the dataset used for training, based on the phase. `ml.Phase` is a string literal with values in `[\"train\", \"valid\"]`. `mlfab.get_data_dir()` returns the data directory, which can be set in a configuration file which lives in `~/.mlfab.yml`. The default configuration file will be written on first run if it doesn't exist yet.\n\n```python\ndef get_dataset(self, phase: mlfab.Phase) -> Dataset[tuple[Tensor, Tensor]]:\n root_dir = mlfab.get_data_dir() / \"mnist\"\n return MNIST(root_dir=root_dir, train=phase == \"train\")\n```\n\n### Compute Loss\n\nEach `mlfab` model should either implement the `forward` function, which should take a batch from the dataset and return the loss, or, if more control is desired, the `get_loss` function can be overridden.\n\n```python\ndef forward(self, x: Tensor) -> Tensor:\n return self.model(x)\n\ndef get_loss(self, batch: tuple[Tensor, Tensor], state: mlfab.State) -> Tensor:\n x_bhw, y_b = batch\n x_bchw = (x_bhw.float() / 255.0).unsqueeze(1)\n yhat_bc = self(x_bchw)\n self.log_step(batch, yhat_bc, state)\n return F.cross_entropy(yhat_bc, y_b)\n```\n\n### Logging\n\nWhen we call `log_step` in the `get_loss` function, it delegates to either `log_train_step` or `log_valid_step`, depending on what `state.phase` is. In this case, on each validation step we log images of the MNIST digits with the labels that our model predicts.\n\n```python\ndef log_valid_step(self, batch: tuple[Tensor, Tensor], output: Tensor, state: mlfab.State) -> None:\n (x_bhw, y_b), yhat_bc = batch, output\n\n def get_label_strings() -> list[str]:\n ypred_b = yhat_bc.argmax(-1)\n return [f\"ytrue={y_b[i]}, ypred={ypred_b[i]}\" for i in range(len(y_b))]\n\n self.log_labeled_images(\"images\", lambda: (x_bhw, get_label_strings()))\n```\n\n### Running\n\nWe can launch a training job using the `launch` class method. The config can be a `Config` object, or it can be the path to a `config.yaml` file located in the same directory as the task file. You can additionally provide the `launcher` argument, which supports training the model across multiple GPUs or nodes.\n\n```python\nif __name__ == \"__main__\":\n MnistClassification.launch(Config(batch_size=16))\n```\n",
"bugtrack_url": null,
"license": null,
"summary": "A collection of core machine learning tools",
"version": "0.2.9",
"project_urls": {
"Homepage": "https://github.com/dpshai/mlfab"
},
"split_keywords": [],
"urls": [
{
"comment_text": "",
"digests": {
"blake2b_256": "14a68dd29a68c36a00508fd752b1bdd3cd21fad0166e9633fb9a94a3f64ef9e3",
"md5": "5e51e824cb322fb4a8d4811885cf7772",
"sha256": "c12b8ef58777239ed3eb4dbde927c05a499165cf9f24f7009e234c38ed021615"
},
"downloads": -1,
"filename": "mlfab-0.2.9-py3-none-any.whl",
"has_sig": false,
"md5_digest": "5e51e824cb322fb4a8d4811885cf7772",
"packagetype": "bdist_wheel",
"python_version": "py3",
"requires_python": ">=3.11",
"size": 208887,
"upload_time": "2024-11-06T23:00:11",
"upload_time_iso_8601": "2024-11-06T23:00:11.216230Z",
"url": "https://files.pythonhosted.org/packages/14/a6/8dd29a68c36a00508fd752b1bdd3cd21fad0166e9633fb9a94a3f64ef9e3/mlfab-0.2.9-py3-none-any.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": "",
"digests": {
"blake2b_256": "6ba3c7ee9791301b0b5232fdc2c6f6929d9550dc901dcfa0b25f7b59b468fe88",
"md5": "72b93bf39992c40e5baaf51adedc75a6",
"sha256": "ff4d694c0893d1c68e83c96a79291cd6771f48e87215d877246b92964efe7bfd"
},
"downloads": -1,
"filename": "mlfab-0.2.9.tar.gz",
"has_sig": false,
"md5_digest": "72b93bf39992c40e5baaf51adedc75a6",
"packagetype": "sdist",
"python_version": "source",
"requires_python": ">=3.11",
"size": 175462,
"upload_time": "2024-11-06T23:00:13",
"upload_time_iso_8601": "2024-11-06T23:00:13.354005Z",
"url": "https://files.pythonhosted.org/packages/6b/a3/c7ee9791301b0b5232fdc2c6f6929d9550dc901dcfa0b25f7b59b468fe88/mlfab-0.2.9.tar.gz",
"yanked": false,
"yanked_reason": null
}
],
"upload_time": "2024-11-06 23:00:13",
"github": true,
"gitlab": false,
"bitbucket": false,
"codeberg": false,
"github_user": "dpshai",
"github_project": "mlfab",
"github_not_found": true,
"lcname": "mlfab"
}
