Aggregate gradients manually over n batches


#1

Dear all,

I have memory limitations so I cannot use large batch size. This has the consequence that my training is unstable. Is there a way to manually aggregate gradients over n-batches of forward steps and manually update my parameters? Am pretty sure there is, so any pointers to code examples extremely appreciated!

Thanks!

PS Any recommendation to manually make training more stable for small batch size is also most welcome.


#2

Are you using Symbolic API or Gluon/autograd?


#3

Thanks, am using Gluon.


#4

This is very straightforward to do with Gluon. You need to set the grad_req in your network Parameter instances to 'add' and manually set the gradient to zero using zero_grad() after each Trainer.step() (see here). To set grad_req to 'add':

for p in net.collect_params().values():
    p.grad_req = 'add'

And similarly call zero_grad() on each parameter after calling Trainer.step(). Remember to modify batch_size argument of trainer.step() accordingly.


#5

thank you very much!!


#6

Hi @safrooze, does this work in the case of data parallelism (multiple copies of the same model in different gpus, with different data sets)? That is, I just need to add all params (in all contexts) grads to ‘add’ and then trainer.step(augmented_batch) after more than one iterations?

Cheers


#7

As long as you follow the steps in multi-GPU tutorial, it should be exactly the same extra steps to set grad_req to 'add', forward/backward multiple times per train.step() followed by zero_grad().


#8

Thanks @safrooze, I tried it and it’s working. This is great help. I went from batch size 32 --> 256, makes a big difference in stability of training. Again, thanks!


#9

Hi, I was using self defined rnn class with ndarray parameters and forward function, with the help of gluon auto grad.

Is it possible to apply the aggregate gradients on this case? Since my model is not a gluon model, so grad_req could not be applied directly.


#10

Hi @ShootingSpace, could you please post your model code? - it will help. From what I see, NDArray has the method (function) attach_grad, that has the option of grad_rec (default: write, but there is add option as well). So in principle, you should be able to perform the same thing.


#11

Thanks for your suggestion! The problem comes out as loss become nan. Do you know if we need to add the loss together over n batches before backward()?
My gru class borrowed from the gluon tutorial

class GRU():
    def __init__(self, vocab_size, num_hidden, seed, ctx=mx.cpu(0)):
        if seed:
            mx.random.seed(2018)

        num_inputs = vocab_size
        num_outputs = vocab_size
        num_hidden = num_hidden

        ########################
        #  Weights connecting the inputs to the hidden layer
        ########################
        self.Wxz = nd.random_normal(shape=(num_inputs,num_hidden), ctx=ctx) * .01
        self.Wxr = nd.random_normal(shape=(num_inputs,num_hidden), ctx=ctx) * .01
        self.Wxh = nd.random_normal(shape=(num_inputs,num_hidden), ctx=ctx) * .01

        ########################
        #  Recurrent weights connecting the hidden layer across time steps
        ########################
        self.Whz = nd.random_normal(shape=(num_hidden,num_hidden), ctx=ctx)* .01
        self.Whr = nd.random_normal(shape=(num_hidden,num_hidden), ctx=ctx)* .01
        self.Whh = nd.random_normal(shape=(num_hidden,num_hidden), ctx=ctx)* .01

        ########################
        #  Bias vector for hidden layer
        ########################
        self.bz = nd.random_normal(shape=num_hidden, ctx=ctx) * .01
        self.br = nd.random_normal(shape=num_hidden, ctx=ctx) * .01
        self.bh = nd.random_normal(shape=num_hidden, ctx=ctx) * .01

        ########################
        # Weights to the output nodes
        ########################
        self.Why = nd.random_normal(shape=(num_hidden,num_outputs), ctx=ctx) * .01
        self.by = nd.random_normal(shape=num_outputs, ctx=ctx) * .01

        self.params = [self.Wxz, self.Wxr, self.Wxh, self.Whz, self.Whr, self.Whh,
                     self.bz, self.br, self.bh, self.Why, self.by]




    def forward(self, inputs, h, temperature=1.0):
        outputs = []
        for X in inputs:
            z = nd.sigmoid(nd.dot(X, self.Wxz) + nd.dot(h, self.Whz) + self.bz)
            r = nd.sigmoid(nd.dot(X, self.Wxr) + nd.dot(h, self.Whr) + self.br)
            g = nd.tanh(nd.dot(X, self.Wxh) + nd.dot(r * h, self.Whh) + self.bh)
            h = z * h + (1 - z) * g

            yhat_linear = nd.dot(h, self.Why) + self.by
            yhat = softmax(yhat_linear, temperature=temperature)
            outputs.append(yhat)
        return (outputs, h)
def cross_entropy(yhat, y):
    return - nd.mean(nd.sum(y * nd.log(yhat), axis=0, exclude=True))

def average_ce_loss(outputs, labels):
    '''Averaging the loss over the sequence'''
    assert(len(outputs) == len(labels))
    total_loss = nd.array([0.], ctx=context)
    for (output, label) in zip(outputs,labels):
        total_loss = total_loss + cross_entropy(output, label)
    return total_loss / len(outputs)

#12

Hi,

am not sure if what I will suggest will work 100%, but it’s easy to give it a try. If I were you, I’d change my model to use gluon directly, from here.

To your example, I think you are missing a line in your model definition, this line:

for param in self.params:
    param.attach_grad(grad_rec = 'add')

I would give it a try with the following modifications (based on the tutorial you followed).

modification in your model:

class GRU():
    def __init__(self, vocab_size, num_hidden, seed, ctx=mx.cpu(0)):
        if seed:
            mx.random.seed(2018)

        num_inputs = vocab_size
        num_outputs = vocab_size
        num_hidden = num_hidden

        ########################
        #  Weights connecting the inputs to the hidden layer
        ########################
        self.Wxz = nd.random_normal(shape=(num_inputs,num_hidden), ctx=ctx) * .01
        self.Wxr = nd.random_normal(shape=(num_inputs,num_hidden), ctx=ctx) * .01
        self.Wxh = nd.random_normal(shape=(num_inputs,num_hidden), ctx=ctx) * .01

        ########################
        #  Recurrent weights connecting the hidden layer across time steps
        ########################
        self.Whz = nd.random_normal(shape=(num_hidden,num_hidden), ctx=ctx)* .01
        self.Whr = nd.random_normal(shape=(num_hidden,num_hidden), ctx=ctx)* .01
        self.Whh = nd.random_normal(shape=(num_hidden,num_hidden), ctx=ctx)* .01

        ########################
        #  Bias vector for hidden layer
        ########################
        self.bz = nd.random_normal(shape=num_hidden, ctx=ctx) * .01
        self.br = nd.random_normal(shape=num_hidden, ctx=ctx) * .01
        self.bh = nd.random_normal(shape=num_hidden, ctx=ctx) * .01

        ########################
        # Weights to the output nodes
        ########################
        self.Why = nd.random_normal(shape=(num_hidden,num_outputs), ctx=ctx) * .01
        self.by = nd.random_normal(shape=num_outputs, ctx=ctx) * .01

        self.params = [self.Wxz, self.Wxr, self.Wxh, self.Whz, self.Whr, self.Whh,
                     self.bz, self.br, self.bh, self.Why, self.by]

        # @@@@@@@@@@@ MODIFICATION HERE @@@@@@@@@@@  
        for param in self.params:
            param.attach_grad(grad_req='add') # This tells mxnet to add the gradients
        # @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@

    def forward(self, inputs, h, temperature=1.0):
        outputs = []
        for X in inputs:
            z = nd.sigmoid(nd.dot(X, self.Wxz) + nd.dot(h, self.Whz) + self.bz)
            r = nd.sigmoid(nd.dot(X, self.Wxr) + nd.dot(h, self.Whr) + self.br)
            g = nd.tanh(nd.dot(X, self.Wxh) + nd.dot(r * h, self.Whh) + self.bh)
            h = z * h + (1 - z) * g

            yhat_linear = nd.dot(h, self.Why) + self.by
            yhat = softmax(yhat_linear, temperature=temperature)
            outputs.append(yhat)
        return (outputs, h)

Now in the example, the SGD function takes place in every iteration, but we need to add manually a “delay_rate”, that will do the update every N iterations, so you have enough aggregated gradients. So I am modifying SGD in the example like this:

delay_rate = 4 # This says to aggregate over 4 batch iterations before updating

# modified SGD that takes the average of the gradients
def SGD(params, lr, _delay_rate):
    for param in params:
        param[:] = param - (lr / _delay_rate) * param.grad

in the initial example, delay_rate is the default behaviour, but if you aggregate gradients over say 4 iterations, you need to divide their magnitude with 4.

Then in the training loop I would replace the line:

SGD(params, learning_rate)

with (assuming you’ve defined the class GRU with the name net somewhere

if (i/delay_rate == 0): # update every delay_rate iterations
    SGD(params, learning_rate,delay_rate)
    # Now manually zero the grads 
    for param in net.params:
         param.zero_grad()

hope this helps. By the way I am a newbie in RNNs, just started learning, so I don’t know if what I say needs modifications for your model.

Cheers


#13

Thanks for your suggest of implementation of the delay sgd,

You are right, I would better transfer to gluon layer.


#14

Hi, I tried this gradients aggregation trick and everything just worked well.
But for the case that I wrote a new customer block where I used softmax() to suppress the learnable weights into range 0~1, I got an error after calling trainer.step() which read:

/softmax-inl.h:267: Check failed: req[0] != kAddTo (3 vs. 3)

Removing the softmax() line in the code can fix the error but it isn’t what I want. Is anything wrong in my code?

def forward(self,x):
      p = nd.softmax(self._rank_p.data(), axis=0)
      .......
      x = nd.dot(x, p)  
      return x

def train():

        net = net()
        for p in net.collect_params().values():
             p.grad_req = 'add'
          .......
        iter_time = 1
        for data, label in train_data:
            with autograd.record():
                output = net(data)
                loss = loss_fun(output, label)
            loss.backward()
            if iter_time % BATCH_UPDATE_PERIOD == 0:
                trainer.step(BATCH_SIZE * BATCH_UPDATE_PERIOD)
                for p in net.collect_params().values():
                    p.zero_grad()
            elif iter_time == len(train_data):
                trainer.step(BATCH_SIZE * (len(train_data) % BATCH_UPDATE_PERIOD))
                for param in net.collect_params().values():
                    param.zero_grad()
            train_loss += nd.mean(loss).asscalar()
            iter_time += 1
......

#15

Hi @Ghostish - You are correct that nd.softmax() doesn’t support grad_req='add'. However I just tested nd.SoftmaxActivation, which basically does the same thing (but is supposed to be deprecated!) and it does support grad_req='add'. So for now, you can use nd.SoftmaxActivation(). I’ve reported the issue here.


#17

Thanks a lot for the solution :grinning: