Do we need to apply exponential moving average to weights during training when we use Adam (or other optimizers)?
My EMA is defined as:
class EMA(object):

def __init__(self, mu):
    self.mu = mu
    self.shadow = {}

def register(self, name, val):
    self.shadow[name] = val.clone()

def __call__(self, name, x):
    assert name in self.shadow
    new_average = (1.0 - self.mu) * x + self.mu * self.shadow[name]
    self.shadow[name] = new_average.clone()
    return new_average

My code is like (during each batch):
        y1, y2 = answer_tpos[:, 0], answer_tpos[:, 1]
        loss1 = self.loss(p1, y1, size_average=True)
        loss2 = self.loss(p2, y2, size_average=True)
        loss = (loss1 + loss2) / 2
        loss.backward()
        self.optimizer.step()

        # update learning rate
        if self.use_scheduler:
            # during warm up stage, use exponential warm up
            if self.step < self.lr_warm_up_num - 1:
                self.scheduler.step()
            # after warm up stage, fix scheduler
            if self.step >= self.lr_warm_up_num - 1 and self.unused:
                self.optimizer.param_groups[0]['initial_lr'] = self.lr
                self.scheduler = optim.lr_scheduler.ExponentialLR(
                    self.optimizer, self.decay)
                for g in self.optimizer.param_groups:
                    g['lr'] = self.lr
                self.unused = False
            # print("Learning rate: {}".format(self.scheduler.get_lr()))
            print("Learning rate: {}".format(
                self.optimizer.param_groups[0]['lr']))

        # exponential moving avarage
        if self.use_ema and self.ema is not None:
            print("Apply ema")
            for name, param in self.model.named_parameters():
                if param.requires_grad:
                    param.data = self.ema(name, param.data)

        # gradient clip
        if self.use_grad_clip:
            torch.nn.utils.clip_grad_norm_(
                self.model.parameters(), self.grad_clip)

And the ema is defined (in main) as:
# set optimizer and scheduler
parameters = filter(lambda p: p.requires_grad, model.parameters())
base_lr = 1.0
optimizer = optim.Adam(
    params=parameters,
    lr=base_lr,
    betas=(args.beta1, args.beta2),
    eps=1e-7,
    weight_decay=3e-7)
cr = args.lr / math.log2(args.lr_warm_up_num)
scheduler = optim.lr_scheduler.LambdaLR(
    optimizer,
    lr_lambda=lambda ee: cr * math.log2(ee + 1)
    if ee < args.lr_warm_up_num else args.lr)

# exponential moving average
ema = EMA(args.decay)
if args.use_ema:
    for name, param in model.named_parameters():
        if param.requires_grad:
            ema.register(name, param.data)

I referred to the following discussion.

  
    
    
    How to apply exponential moving average decay for variables? 
  
  
    I am reading following paper. And it uses EMA decay for variables. 
BI-DIRECTIONAL ATTENTION FLOW FOR MACHINE COMPREHENSION 

During training, the moving averages of all weights of the model are 
maintained with the exponential decay rate of 0.999. 

They use TensorFlow and I found the related code of EMA. 

 
In PyTorch, how do I apply EMA to Variables? In TensorFlow, there is tf.train.ExponentialMovingAverage class. 
https://www.tensorflow.org/versions/r0.12/api_docs/python/train/moving_averag…
  


My problem is:

whether my implementation is correct? It seems that after I applied ema to weights, my training performance is not good.
seems that some optimizers use exponential moving average on gradients. So if they used EMA for gradients, do we need to apply EMA for weights again after optimizer?