pytorch中使用Lion优化器

🦁Lion，谷歌大脑在 Pytorch 中使用遗传算法发现的新优化器，据称比 Adam(w) 更好。

Github

安装

$ pip install lion-pytorch

使用

# toy model

import torch
from torch import nn

model = nn.Linear(10, 1)

# import Lion and instantiate with parameters

from lion_pytorch import Lion

opt = Lion(model.parameters(), lr=1e-4, weight_decay=1e-2)

# forward and backwards

loss = model(torch.randn(10))
loss.backward()

# optimizer step

opt.step()
opt.zero_grad()