torch_optimizer_accsgd_quickstart_single_training_step.py

1import torch
2import torch_optimizer as optim
3
4# 1. Define a simple model and input
5model = torch.nn.Linear(10, 1)
6input = torch.randn(1, 10)
7target = torch.randn(1, 1)
8
9# 2. Initialize the optimizer from torch-optimizer
10# You can replace AccSGD with any other optimizer like RAdam, DiffGrad, etc.
11optimizer = optim.AccSGD(
12    model.parameters(),
13    lr=1e-3,
14    weight_decay=0,
15)
16
17# 3. Standard PyTorch training step
18optimizer.zero_grad()
19loss_fn = torch.nn.MSELoss()
20output = model(input)
21loss = loss_fn(output, target)
22loss.backward()
23
24# 4. Perform the optimization step
25optimizer.step()