在Pytorch中实现自定义学习率调度器？

import torch.optim.lr_scheduler as lr_scheduler

class MyScheduler(lr_scheduler._LRScheduler):
    def __init__(self, optimizer, ...):
        super(MyScheduler, self).__init__(optimizer)
        # 定义初始化参数
        
    def get_lr(self):
        # 实现学习率调度逻辑
        # 返回一个列表，其中包含了每个参数组的学习率

model = ...
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
scheduler = MyScheduler(optimizer, ...)
...
for epoch in range(num_epochs):
    train(...)
    scheduler.step()  # 每个epoch结束后调用