Memory Issues

[mxtsai]

Hi Luke,

Thank you for the awesome work. I tried running EfficientNet-B0 on my GTX 1070 (8GB RAM) with an input batch of dimension [44x1x256x256] (single channel image) and I am running into 'CUDA out of memory' (with the model in 'training' mode).

I tried running another implementation and wasn't getting this issue, and after digging in the code, it seems as if the implementation for MBConv (or the re-iteration of MBConv) was too memory hungry.

I really like your implementation of EfficientNet and if I did have more time, I would definitely have a deeper dive into your code. At the mean time, if possible, could you help me check this issue out (maybe it'll speed up training in the future?) ? Thank you!

[lukemelas]

Hello, thanks for filing the issue. I'll look into it this weekend.

[jlevy44]

I think I'm running into a similar issue as well.

[crazyblacker]

I meet the same issue when I exchange backbone with B3. Model size become smaller but out of memory on my 2080Ti(11G).

[dami23]

Hi @lukemelas. Thank you for your awesome work. I find the newly released version requires more GPU memory than the original one. The same project use B3 model can be run with the older version, but after the upgrade, this code cannot be run anymore on the same workstation and with error"RuntimeError: CUDA out of memory". Could you please give me any advices to face this problem? Thank you!!

[lukemelas]

Hi @dami23 , that's strange because nothing about the B3 model changed in any way during the update (apart from the addition of drop_connect). In fact, you can see all the changes that were made here: 9a0d45f

Can you try measuring memory usage with the two versions to confirm the issue?

[xbpgithub]

Hello，I also meet the similar memory issue. EfficinetNet-b3 consumes more memory than Resnet50 for FPN, is it normal? Can someone provide some help?

[twmht]

any update on this?

[seilna]

Similar problems with @xbpgithub, and it seems that in MBConvBlock(), relu_fn(x) and self._bn are memory bottleneck. Replacing them with Inplace BatchNorm was helpful to significantly reduce memory usage, but until now it does not support swish activation function (it only supports invertible activation functions like leaky_relu).
(I'm not sure it is right direction to focus on batchnorm to reduce memory usage, but hope it can be helpful)

[seilna]

I also observed that in relu_fn(), using custom op for swish activation function (instead of return x * torch.sigmoid(x)) was helpful to reduce memory usage (up to 30%) as follows:

sigmoid = torch.nn.Sigmoid()
class Swish(torch.autograd.Function):
    @staticmethod
    def forward(ctx, i):
        result = i * sigmoid(i)
        ctx.save_for_backward(i)
        return result

    @staticmethod
    def backward(ctx, grad_output):
        i = ctx.saved_variables[0]
        sigmoid_i = sigmoid(i)
        return grad_output * (sigmoid_i * (1 + i * (1 - sigmoid_i)))

swish = Swish.apply

class Swish_module(nn.Module):
    def forward(self, x):
        return swish(x)

swish_layer = Swish_module()
def relu_fn(x):
    """ Swish activation function """
    # return x * torch.sigmoid(x)
    return swish_layer(x)

[jackytu256]

@seilna
perhaps, the h-swish (provided by MobilenetV3) may be another option to deal with this issue?

class h_sigmoid(nn.Module):
    def __init__(self, inplace=True):
        super(h_sigmoid, self).__init__()
        self.relu = nn.ReLU6(inplace=inplace)

    def forward(self, x):
        return self.relu(x + 3) / 6


class h_swish(nn.Module):
    def __init__(self, inplace=True):
        super(h_swish, self).__init__()
        self.sigmoid = h_sigmoid(inplace=inplace)

    def forward(self, x):
        return x * self.sigmoid(x)

[tmabraham]

Any update on this? @DrHB mentioned this over here and it would be nice if this repository was updated with this fix.

[lukemelas]

The latest version of the repo (released a week ago) should fix this issue!

A memory-efficient swish function is now used by default. However, this swish function is incompatible with model exporting (ONNX), so you also have the option of using model.set_swish(memory_efficient=False).

[lucasjinreal]

@lukemelas Is that a way to make ONNX also memory efficient?