nixoptimiz 0.1.0

Nix Optimizer in PyTorch

Nix

Nix is an simple optimizer for optimize complex models. Nix是一个用于优化复杂模型的简单优化器。

Usage 使用

Here is example code 以下是示例代码：

from nixoptimiz import NixOptimizer
import torch
from torch import nn

# Example Model
model = nn.Sequential(
    nn.Linear(16, 512),
    nn.ReLU(),
    nn.Linear(512, 16),
    nn.Sigmoid()
).cuda()

# Example Data
inputs = torch.randn(64, 16).cuda()
outputs = torch.rand(64, 16).cuda()

optim = NixOptimizer(model.parameters(), lr=0.01, muon_mode="kernel")
criterion = nn.BCELoss()

for idx in range(0, 5001):
    loss = criterion(model(inputs), outputs)
    optim.zero_grad()
    loss.backward()
    optim.step()
    if idx % 100 == 0:
        print(f"Train step: {idx}, Loss: {float(loss.cpu())}")