python-3.x - pytorch で手書きのオプティマイザが機能しない

Question

私はこのアルゴリズムに従ってオプティマイザーを書いてみました: アルゴリズム 私がアルゴリズムを取った記事 (5 ページ): https://arxiv.org/pdf/2106.02720.pdf

しかし、私の解決策は何も最適化していません。修正しようとしたら、d_pがゼロだけで構成されていることがわかりましたが、理由がわかりませんでした(

import torch
from torch.optim import optimizer
class BAMSGD(optimizer.Optimizer):
  def __init__(self, params, lr=float(1e-2), weight_decay=0, gamma=0.9):
    default = dict(lr=lr, weight_decay=weight_decay, gamma=gamma)
    super(BAMSGD, self).__init__(params, default)

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

  def step(self, closure=None):

    loss = None
    if closure is not None:
      loss=closure()

    for group in self.param_groups:
      weight_decay = group["weight_decay"]
      learning_rate = group["lr"]
      gamma = group["gamma"]
      for p in range(len(group["params"])):
        t = p
        p = group["params"][p]
        param_state = self.state[p]
        if "wt_ag" not in param_state:
          param_state["wt_ag"] = copy.deepcopy(p.data)
        if "wt" not in param_state:
          param_state["wt"] = copy.deepcopy(p.data)
        beta_t = 1 + t/6
        gamma_t = gamma * (t+1)
        b = 200 # at the moment it's just a random number
        wt_ag = param_state["wt_ag"]
        wt = param_state["wt"]
        if p.grad is None:
            continue
        d_p = p.grad.data
        d_p.add_(1e-2, p.data)
        '''if weight_decay != 0:
            d_p.add_(weight_decay, p.data)'''
        # buffer.mul_(1-pow(beta_t, -1)) #
        # buffer2.mul_(pow(beta_t, -1)) #
        # buffer.add_(buffer2) #
        wt.add(-(gamma * d_p)) ##
        wt_md = pow(beta_t, -1) * wt + (1-pow(beta_t, -1)) * wt_ag
        ww = wt - gamma * d_p
        w_t1 = min(1, b/(wt.norm())) * wt
        w_t1ag = pow(beta_t, -1) * w_t1 + (1 - pow(beta_t, -1)) * wt_ag
        # self.__setstate__(w_t1ag, w_t1)
        p.data.add_(-group['lr'], w_t1)
    return loss

私を助けてください、私は最適化の初心者であり、コードで多くの間違いを犯した可能性があります(