给出一个torch自定义优化器的程序

下面是一个使用 PyTorch 自定义优化器的程序：

import torch.optim as optim

class MyOptimizer(optim.Optimizer):
    def __init__(self, params, lr=0.01):
        defaults = dict(lr=lr)
        super(MyOptimizer, self).__init__(params, defaults)

    def __setstate__(self, state):
        super(MyOptimizer, self).__setstate__(state)

    def step(self, closure=None):
        loss = None
        if closure is not None:
            loss = closure()

        for group in self.param_groups:
            lr = group['lr']

            for p in group['params']:
                if p.grad is None:
                    continue
                grad = p.grad.data
                p.data.add_(-lr, grad)

        return loss

在这个例子中，我们定义了一个名为 `MyOptimizer` 的自定义优化器，它继承自 `torch.optim.Optimizer` 类。在初始化函数中，我们设置了优化器的学习率，然后调用父类的 `__init__` 函数。在 `step` 函数中，我们遍历参数组并更新参数。注意，我们使用了张量的 `add_` 函数来原地更新张量的值，这可以减少内存使用和计算时间。

使用这个自定义优化器的方法与使用其他 PyTorch 优化器的方法相同，例如：

optimizer = MyOptimizer(model.parameters(), lr=0.1)
optimizer.zero_grad()
loss_fn(model(input), target).backward()
optimizer.step()