[Bug] backward after optimzer.step()? about mmrazor HOT 14 CLOSED

Forcing optimize_config to None may not be a good idea, for example, it's not easy to adapt to fp16 training when setting optimizer_config to None.

Why not just comment out all the .backward and .step in https://github.com/open-mmlab/mmrazor/blob/master/mmrazor/models/algorithms/autoslim.py#L226

I tried and it works.

Thank you for your issue.
It does work to comment out all the .backward and .step. And we use these here just to save GPU memory. It is hacky as it's not easy to adapt to fp16 training. But just using OptimizerHook may lead to cuda out of memory when training on bigger dataset, such as Imagenet. We are still working on it.

from mmrazor.

Maybe I didn't make it clear.
code 1:

loss1 = forward_loss()
loss2 = forward_loss()
loss3 = forward_loss()
loss = loss1 + loss2 + loss3
loss.backward()
optimizor.step()

code2:

loss1 = forward_loss()
loss1.backward()
loss2 = forward_loss()
loss2.backward()
loss3 = forward_loss()
loss3.backward()
optimizor.step()

The memory cost of code1 is 3 times as much as that of code2.
If we use OptimizerHook, the effect is similar to code1.

from mmrazor.

I ran into this issue just now. And I solve it by setting optimizer_config = None as

from mmrazor.

Forcing optimize_config to None may not be a good idea, for example, it's not easy to adapt to fp16 training when setting optimizer_config to None.

Why not just comment out all the .backward and .step in https://github.com/open-mmlab/mmrazor/blob/master/mmrazor/models/algorithms/autoslim.py#L226

I tried and it works.

from mmrazor.

@HIT-cwh

I am wondering why the current implementation would result OOM, I reimplement with SlimConv2D from https://github.com/JiahuiYu/slimmable_networks/blob/master/models/slimmable_ops.py#L25.

And there is no OOM occur.

from mmrazor.

@HIT-cwh

from mmrazor.

I am wondering why the current implementation would result OOM, I reimplement with SlimConv2D from https://github.com/JiahuiYu/slimmable_networks/blob/master/models/slimmable_ops.py#L25.

And there is no OOM occur.

It is better to make a new issue to discuss this question. @twmht

from mmrazor.

@HIT-cwh

I am wondering why the current implementation would result OOM, I reimplement with SlimConv2D from https://github.com/JiahuiYu/slimmable_networks/blob/master/models/slimmable_ops.py#L25.

And there is no OOM occur.

Look here. In the official code, the gradient of the loss is computed immediately for each width_mult. That is, loss.backward() is executed immediately after forward_loss. So there is no OOM occur.
But in our implementation, if we use OptimizerHook, the pseudo code will be roughly as follows:

loss = 0.   
for i, subnet in enumerate(self.channel_cfg):   
    ...   
    model_loss, _ = self._parse_losses(model_losses)  
    loss += model_loss  
loss.backward()  

It will use much more memory.

from mmrazor.

@HIT-cwh

But you can also do in the same way with mmrazor, where is the OOM from if you do in the same way ?

from mmrazor.

Interesting. To the best of my knowledge, I though both use the same memory. Thank you for pointing out that.

from mmrazor.

Interesting. To the best of my knowledge, I though both use the same memory. Thank you for pointing out that.

I'm sorry to bother you again. To some degree, the above code2 has the same effect as grad-cumulating. So, I think the implementation like code2 may save much memory. If it turns out that I'm wrong, please let me know. Thanks.

from mmrazor.

Yup. I think you are right, since we need to trace the autograd graph when training.

from mmrazor.

@HIT-cwh

Well, I found out it also works by commenting out all the backward and step when training object detection, No OOM happens.

It's weird since the consuming memory is almost the same as when doing normal object detection training.

My pytorch version is 1.9.

What is your batch size and number of gpus when training with ImageNet?

from mmrazor.

Samples_per_gpu=128 and GPUS=8

from mmrazor.