Skip to content

Instantly share code, notes, and snippets.

@MiXaiLL76
Created January 13, 2022 15:47
Show Gist options
  • Save MiXaiLL76/a7844224e1ac891b9f16ddcf0b2e533b to your computer and use it in GitHub Desktop.
Save MiXaiLL76/a7844224e1ac891b9f16ddcf0b2e533b to your computer and use it in GitHub Desktop.
import matplotlib.pyplot as plt
import torch
from torch.optim import SGD
from torch.optim.lr_scheduler import ExponentialLR, CosineAnnealingWarmRestarts, CosineAnnealingLR
model = [torch.nn.Parameter(torch.randn(2, 2, requires_grad=True))]
GPU_COUNT = 1
BASE_LR = 0.02 # 2e-2
# BASE_LR = 0.2
bs = (GPU_COUNT * 2)
optimizer_lr = (BASE_LR * bs / 16)
print(f"{optimizer_lr=}")
epoches = 100
step_per_epoch = 70815 // 32
step_per_epoch = 1
optimizer = SGD(model, optimizer_lr)
scheduler1 = ExponentialLR(optimizer, gamma=0.9)
scheduler2 = CosineAnnealingWarmRestarts(optimizer, T_0=epoches//4, T_mult=1, eta_min=1e-4)
# scheduler3 = CosineAnnealingLR(optimizer, T_max=epoches, eta_min=1e-4)
lr = []
s = scheduler1
s = scheduler2
# s = scheduler
for epoch in range(epoches):
for _ in range(step_per_epoch):
optimizer.step()
lr.append(round(s.get_last_lr()[0], 5))
s.step(epoch)
print(max(lr), min(lr))
plt.plot(lr)
plt.show()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment