tensorneko

KO | EN

gitlite — search

#python #java #python3 #arduino #golang #machine-learning #rust #html #flask #javascript #seismology #nodejs

★ 10 Open GitHub ↗

Tensor Neural Engine Kompanion. An util library based on PyTorch and PyTorch Lightning.

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

<h1 style="text-align: center">TensorNeko</h1>

Tensor Neural Engine Kompanion. An util library based on PyTorch and PyTorch Lightning.

## Install

The tensorneko requires pytorch and pytorch-lightning (optional), and you can install it with below command.

```shell
pip install tensorneko  # for PyTorch only
pip install tensorneko[lightning]  # for PyTorch and Lightning
```

To use the library without PyTorch and PyTorch Lightning, you can install the util library (support Python 3.7 ~ 3.14 with limited features) with following command.
```shell
pip install tensorneko_util
```

Some cpu bound functions are implemented by rust-based `pyo3`, and you can install the optimized version with below command.
```shell
pip install tensorneko_lib
```

Some CLI tools are provided in the `tensorneko_tool` package, and you can install it with below command.
```shell
pipx install tensorneko_tool  # or `pip install tensorneko_tool`
```

Then you can use the CLI tools `tensorneko` in the terminal.

## Layers, Modules and Architectures

Build an MLP with linear layers. The activation and normalization will be placed in the hidden layers.

784 -> 1024 -> 512 -> 10

```python
import tensorneko as neko
import torch.nn

mlp = neko.module.MLP(
    neurons=[784, 1024, 512, 10],
    build_activation=torch.nn.ReLU,
    build_normalization=[
        lambda: torch.nn.BatchNorm1d(1024),
        lambda: torch.nn.BatchNorm1d(512)
    ],
    dropout_rate=0.5
)
```

Build a Conv2d with activation and normalization.

```python
import tensorneko as neko
import torch.nn

conv2d = neko.layer.Conv2d(
    in_channels=256,
    out_channels=1024,
    kernel_size=(3, 3),
    padding=(1, 1),
    build_activation=torch.nn.ReLU,
    build_normalization=lambda: torch.nn.BatchNorm2d(256),
    normalization_after_activation=False
)
```

#### All architectures, modules and layers

Layers:

- `Aggregation`
- `Concatenate`
- `Conv`, `Conv1d`, `Conv2d`, `Conv3d`
- `GaussianNoise`
- `ImageAttention`, `SeqAttention`
- `MaskedConv2d`, `MaskedConv2dA`, `MaskedConv2dB`
- `Linear`
- `Log`
- `PatchEmbedding2d`
- `PositionalEmbedding`
- `Reshape`
- `Stack`
- `VectorQuantizer`

Modules:

- `DenseBlock`
- `InceptionModule`
- `MLP`
- `ResidualBlock` and `ResidualModule`
- `AttentionModule`, `TransformerEncoderBlock` and `TransformerEncoder`
- `GatedConv`

Architectures:
- `AutoEncoder`
- `GAN`
- `WGAN`
- `VQVAE`

## Neko modules

All `tensorneko.layer` and `tensorneko.module` are `NekoModule`. They can be used in 
[fn.py](https://github.com/kachayev/fn.py) pipe operation.

```python
from tensorneko.layer import Linear
from torch.nn import ReLU
import torch

linear0 = Linear(16, 128, build_activation=ReLU)
linear1 = Linear(128, 1)

f = linear0 >> linear1
print(f(torch.rand(16)).shape)
# torch.Size([1])
```

## IO

Easily load and save different modal data.

```python
import tensorneko as neko
from tensorneko.io import json_data
from typing import List

# read video (Temporal, Channel, Height, Width)
video_tensor, audio_tensor, video_info = neko.io.read.video("path/to/video.mp4")
# write video
neko.io.write.video("path/to/video.mp4", 
    video_tensor, video_info.video_fps,
    audio_tensor, video_info.audio_fps
)

# read audio (Channel, Temporal)
audio_tensor, sample_rate = neko.io.read.audio("path/to/audio.wav")
# write audio
neko.io.write.audio("path/to/audio.wav", audio_tensor, sample_rate)

# read image (Channel, Height, Width) with float value in range [0, 1]
image_tensor = neko.io.read.image("path/to/image.png")
# write image
neko.io.write.image("path/to/image.png", image_tensor)
neko.io.write.image("path/to/image.jpg", image_tensor)

# read plain text
text_string = neko.io.read.text("path/to/text.txt")
# write plain text
neko.io.write.text("path/to/text.txt", text_string)

# read json as python dict or list
json_dict = neko.io.read.json("path/to/json.json")
# read json as an object
@json_data
class JsonData:
    x: int
    y: int

json_obj: List[JsonData] = neko.io.read.json("path/to/json.json", cls=List[JsonData])
# write json from python dict/list or json_data decorated object
neko.io.write.json("path/to/json.json", json_dict)
neko.io.write.json("path/to/json.json", json_obj)
```

Besides, the read/write for `mat` and `pickle` files is also supported.

## Preprocessing

```python
import tensorneko as neko

# A video tensor with (120, 3, 720, 1280)
video = neko.io.read.video("example/video.mp4").video
# Get a resized tensor with (120, 3, 256, 256)
resized_video = neko.preprocess.resize_video(video, (256, 256))
```

#### All preprocessing utils

- `resize_video`
- `resize_image`
- `padding_video`
- `padding_audio`
- `crop_with_padding`
- `frames2video`

if `ffmpeg` is available, you can use below ffmpeg wrappers.

- `video2frames`
- `merge_video_audio`
- `resample_video_fps`
- `mp32wav`

## Visualization

### Variable Web Watcher
Start a web server to watch the variable status when the program (e.g. training, inference, data preprocessing) is running.
```python
import time
from tensorneko.visualization.watcher import *
data_list = ... # a list of data
def preprocessing(d): ...

# initialize the components
pb = ProgressBar("Processing", total=len(data_list))
logger = Logger("Log message")
var = Variable("Some Value", 0)
line_chart = LineChart("Line Chart", x_label="x", y_label="y")
view = View("Data preprocessing").add_all()

t0 = time.time()
# open server when the code block in running.
with Server(view, port=8000):
    for i, data in enumerate(data_list):
        preprocessing(data) # do some processing here
        
        x = time.time() - t0  # time since the start of the program
        y = i # processed number of data
        line_chart.add(x, y)  # add to the line chart
        logger.log("Some messages")  # log messages to the server
        var.value = ...  # keep tracking a variable
        pb.add(1)  # update the progress bar by add 1
```
When the script is running, go to `127.0.0.1:8000` to keep tracking the status.

### Tensorboard Server

Simply run tensorboard server in Python script.
```python
import tensorneko as neko

with neko.visualization.tensorboard.Server(port=6006):
    trainer.fit(model, dm)
```

### Matplotlib wrappers
Display an image of (C, H, W) shape by `plt.imshow` wrapper.
```python
import tensorneko as neko
import matplotlib.pyplot as plt

image_tensor = ...  # an image tensor with shape (C, H, W)
neko.visualization.matplotlib.imshow(image_tensor)
plt.show()
```

### Predefined colors
Several aesthetic colors are predefined.

```python
import tensorneko as neko
import matplotlib.pyplot as plt

# use with matplotlib
plt.plot(..., color=neko.visualization.Colors.RED)

# the palette for seaborn is also available
from tensorneko_util.visualization.seaborn import palette
import seaborn as sns
sns.set_palette(palette)
```

## Neko Model

Build and train a simple model for classifying MNIST with MLP.

```python
from typing import Optional, Union, Sequence, Dict, List

import torch.nn
from torch import Tensor
from torch.optim import Adam
from torchmetrics import Accuracy
from lightning.pytorch.callbacks import ModelCheckpoint

import tensorneko as neko
from tensorneko.util import get_activation, get_loss

class MnistClassifier(neko.NekoModel):

def __init__(self, name: str, mlp_neurons: List[int], activation: str, dropout_rate: float, loss: str,
        learning_rate: float, weight_decay: float
    ):
        super().__init__(name)
        self.weight_decay = weight_decay
        self.learning_rate = learning_rate

self.flatten = torch.nn.Flatten()
        self.mlp = neko.module.MLP(
            neurons=mlp_neurons,
            build_activation=get_activation(activation),
            dropout_rate=dropout_rate
        )
        self.loss_func = get_loss(loss)()
        self.acc_func = Accuracy()

def forward(self, x):
        # (batch, 28, 28)
        x = self.flatten(x)
        # (batch, 768)
        x = self.mlp(x)
        # (batch, 10)
        return x

def training_step(self, batch: Optional[Union[Tensor, Sequence[Tensor]]] = None, batch_idx: Optional[int] = None,
        optimizer_idx: Optional[int] = None, hiddens: Optional[Tensor] = None
    ) -> Dict[str, Tensor]:
        x, y = batch
        logit = self(x)
        prob = logit.sigmoid()
        loss = self.loss_func(logit, y)
        acc = self.acc_func(prob.max(dim=1)[1], y)
        return {"loss": loss, "acc": acc}

def validation_step(self, batch: Optional[Union[Tensor, Sequence[Tensor]]] = None, batch_idx: Optional[int] = None,
        dataloader_idx: Optional[int] = None
    ) -> Dict[str, Tensor]:
        x, y = batch
        logit = self(x)
        prob = logit.sigmoid()
        loss = self.loss_func(logit, y)
        acc = self.acc_func(prob.max(dim=1)[1], y)
        return {"loss": loss, "acc": acc}

def configure_optimizers(self):
        optimizer = Adam(self.parameters(), lr=self.learning_rate, betas=(0.5, 0.9), weight_decay=self.weight_decay)
        return {
            "optimizer": optimizer
        }

model = MnistClassifier("mnist_mlp_classifier", [784, 1024, 512, 10], "ReLU", 0.5, "CrossEntropyLoss", 1e-4, 1e-4)

dm = ...  # The MNIST datamodule from PyTorch Lightning

trainer = neko.NekoTrainer(log_every_n_steps=100, gpus=1, logger=model.name, precision=32,
    callbacks=[ModelCheckpoint(dirpath="./ckpt",
        save_last=True, filename=model.name + "-{epoch}-{val_acc:.3f}", monitor="val_acc", mode="max"
    )])

trainer.fit(model, dm)
```

## Callbacks

Some simple but useful pytorch-lightning callbacks are provided.

- `DisplayMetricsCallback`
- `EarlyStoppingLR`: Early stop training when learning rate reaches threshold.

## Notebook Helpers
Here are some helper functions to better interact with Jupyter Notebook.
```python
import tensorneko as neko
# display a video
neko.notebook.display.video("path/to/video.mp4")
# display an audio
neko.notebook.display.audio("path/to/audio.wav")
# display a code file
neko.notebook.display.code("path/to/code.java")
```

## Debug Tools

Get the default values from `ArgumentParser` args. It's convenient to use this in the notebook.
```python
from argparse import ArgumentParser
from tensorneko.debug import get_parser_default_args

parser = ArgumentParser()
parser.add_argument("integers", type=int, nargs="+", default=[1, 2, 3])
parser.add_argument("--sum", dest="accumulate", action="store_const", const=sum, default=max)
args = get_parser_default_args(parser)

print(args.integers)  # [1, 2, 3]
print(args.accumulate)  # <function sum at ...>
```

## Evaluation

Some metrics function for evaluation are provided.

- `iou_1d`
- `iou_2d`
- `psnr_video`
- `psnr_image`
- `ssim_video`
- `ssim_image`

## Message (Access to other services)

### Gotify

Send a message to the Gotify server.

The title, URL and APP_TOKEN is the environment variable `GOTIFY_TITLE`, `GOTIFY_URL` and `GOTIFY_TOKEN`, or overwritten 
in the function arguments.

```python
from tensorneko.msg import gotify
gotify.push("This is a test message", "<URL>", "<APP_TOKEN>")
# then the message will be sent to the Gotify server.
# title = "<HOST_NAME>", message = "This is a test message", priority = 0
```

### Postgres

Require the `psycopg` package. Provide one single function to execute one SQL query with a temp connection.

The database URL is the environment variable `DB_URL`, or overwritten in the function arguments.
```python
from tensorneko.msg import postgres
result = postgres.execute("<SQL>", "<DB_URL>")
# also async version is provided
result = await postgres.execute_async("<SQL>", "<DB_URL>")
```

## Utilities

### Misc functions

`__`: The arguments to pipe operator. (Inspired from [fn.py](https://github.com/kachayev/fn.py))
```python
from tensorneko.util import __, _
result = __(20) >> (_ + 1) >> (_ * 2) >> __.get
print(result)
# 42
```

`Seq` and `Stream`: A collection wrapper for method chaining with concurrent supporting.
```python
from tensorneko.util import Seq, Stream, _
from tensorneko_util.backend.parallel import ParallelType
# using method chaining
seq = Seq.of(1, 2, 3).map(_ + 1).filter(_ % 2 == 0).map(_ * 2).take(2).to_list()
# return [4, 8]

# using bit shift operator to chain the sequence
seq = Seq.of(1, 2, 3) << Seq.of(2, 3, 4) << [3, 4, 5]
# return Seq(1, 2, 3, 2, 3, 4, 3, 4, 5)

# run concurrent with `for_each` for Stream
if __name__ == '__main__':
    Stream.of(1, 2, 3, 4).for_each(print, progress_bar=True, parallel_type=ParallelType.PROCESS)
```

`Option`: A monad for dealing with data.
```python
from tensorneko.util import return_option

@return_option
def get_data():
    if some_condition:
        return 1
    else:
        return None

def process_data(n: int):
    if condition(n):
        return n
    else:
        return None

data = get_data()
data = data.map(process_data).get_or_else(-1)  # if the response is None, return -1
```

`Eval`: A monad for lazy evaluation.
```python
from tensorneko.util import Eval

@Eval.always
def call_by_name_var():
    return 42

@Eval.later
def call_by_need_var():
    return 43

@Eval.now
def call_by_value_var():
    return 44

print(call_by_name_var.value)  # 42
```

### Reactive
This library provides event bus based reactive tools. The API integrates the Python type annotation syntax.

```python
# useful decorators for default event bus
from tensorneko.util import subscribe
# Event base type
from tensorneko.util import Event, EventBus

class LogEvent(Event):
    def __init__(self, message: str):
        self.message = message

# the event argument should be annotated correctly
@subscribe # run in the main thread
def log_information(event: LogEvent):
    print(event.message)

@subscribe.thread # run in a new thread
def log_information_thread(event: LogEvent):
    print(event.message, "in another thread")

@subscribe.coro # run with async
async def log_information_async(event: LogEvent):
    print(event.message, "async")

@subscribe.process # run in a new process
def log_information_process(event: LogEvent):
    print(event.message, "in a new process")

if __name__ == '__main__':
    # emit an event, and then the event handler will be invoked
    # The sequential order is not guaranteed
    LogEvent("Hello world!")
    EventBus.default.wait()  # it's not blocking, need to call wait manually before exit.
    # one possible output:
    # Hello world! in another thread
    # Hello world! async
    # Hello world!
    # Hello world! in a new process
```

### Multiple Dispatch

`dispatch`: Multi-dispatch implementation for Python.

To my knowledge, 3 popular multi-dispatch libraries still have critical limitations. 
[plum](https://github.com/wesselb/plum) doesn't support static methods, 
[mutipledispatch](https://github.com/mrocklin/multipledispatch) doesn't support Python type annotation syntax and 
[multimethod](https://github.com/coady/multimethod) doesn't support default argument. TensorNeko can do it all.

```python
from tensorneko.util import dispatch

class DispatchExample:

@staticmethod
    @dispatch
    def go() -> None:
        print("Go0")

@staticmethod
    @dispatch
    def go(x: int) -> None:
        print("Go1")

@staticmethod
    @dispatch
    def go(x: float, y: float = 1.0) -> None:
        print("Go2")

@dispatch
def come(x: int) -> str:
    return "Come1"

@dispatch.of(str)
def come(x) -> str:
    return "Come2"
```

### Miscellaneous

`StringGetter`: Get PyTorch class from string.
```python
import tensorneko as neko
activation = neko.util.get_activation("leakyRelu")()
```

`Seed`: The universal seed for `numpy`, `torch` and Python `random`.
```python
from tensorneko.util import Seed
from torch.utils.data import DataLoader

# set seed to 42 for all numpy, torch and python random
Seed.set(42)

# Apply seed to parallel workers of DataLoader
DataLoader(
    train_dataset,
    batch_size=batch_size,
    num_workers=num_workers,
    worker_init_fn=Seed.get_loader_worker_init(),
    generator=Seed.get_torch_generator()
)
```

`Timer`: A timer for measuring the time.
```python
from tensorneko.util import Timer
import time

# use as a context manager with single time
with Timer():
    time.sleep(1)

# use as a context manager with multiple segments
with Timer() as t:
    time.sleep(1)
    t.time("sleep A")
    time.sleep(1)
    t.time("sleep B")
    time.sleep(1)

# use as a decorator
@Timer()
def f():
    time.sleep(1)
    print("f")
```

`Singleton`: A decorator to make a class as a singleton. Inspired from Scala/Kotlin.
```python
from tensorneko.util import Singleton

@Singleton
class MyObject:
    def __init__(self):
        self.value = 0

def add(self, value):
        self.value += value
        return self.value

print(MyObject.value)  # 0
MyObject.add(1)
print(MyObject.value)  # 1
```

Besides, many miscellaneous functions are also provided.

Functions list (in `tensorneko_util`):
- `generate_inf_seq`
- `compose`
- `listdir`
- `with_printed`
- `ifelse`
- `dict_add`
- `as_list`
- `identity`
- `list_to_dict`
- `get_tensorneko_util_path`

Functions list (in `tensorneko`):
- `reduce_dict_by`
- `summarize_dict_by`
- `with_printed_shape`
- `is_bad_num`
- `count_parameters`

## TensorNeko Tools

Some CLI tools are provided in the `tensorneko_tool` package. After installing `tensorneko_tool`, you can use the `tensorneko` command in the terminal.

Global flags:

- `tensorneko --version`: print the installed CLI version
- `tensorneko --banner`: show the TensorNeko banner
- `tensorneko --quiet ...`: suppress normal terminal output

Available command families:

- `gotify`: send a Gotify notification or watch a process and notify when it exits
- `dep_check`: compare installed packages with a `requirements.txt` file
- `openai`: test, list, and chat with an OpenAI-compatible endpoint

### Gotify

Set `GOTIFY_URL` and `GOTIFY_TOKEN` in the environment, or pass them with `--url` and `--token`.

```shell
tensorneko gotify "Script finished!"
tensorneko gotify --watch python "Training finished!"
```

### Dependency checker

Compare the current environment with a requirements file.

```shell
tensorneko dep_check -r requirements.txt
tensorneko dep_check -r requirements.txt --overwrite
```

### OpenAI-compatible CLI

Use `openai test` to validate an endpoint, `openai list` to inspect available models, and `openai chat` to send a quick request.

```shell
tensorneko openai test --endpoint https://api.openai.com/v1 --key "$OPENAI_API_KEY"
tensorneko openai list --endpoint https://api.openai.com/v1 --key "$OPENAI_API_KEY"
tensorneko openai chat "Say hello in one sentence." --endpoint https://api.openai.com/v1 --key "$OPENAI_API_KEY"
```

The `openai test` command also supports `--mode network|auth|models|probe|all`, and `openai chat` / `openai list` support `--json` for scripting.

tensorneko

Explore Similar Repositories

Was this content helpful?

Recent Feedback

Download README