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())}")