This paper explored the trade-off between stability and efficiency in Convolution-BatchNorm blocks which are popular convolution neural networks：Deploy mode is efficient but suffers from training instability; Eval mode is widely used in transfer learning but lacks efficiency. Based on detailed analysis, the paper proposed a Tune mode, which is stable and efficient, to bridge the gap between Eval mode and Deploy mode. Numerous experiments conducted in various tasks have verified the effectiveness of the proposed Tune mode.

In this paper, the authors introduce the Tune mode of ConvBN blocks, which can save memory and time costs during transfer learning. Experiments show that when transferring the backbone with ConvBN to downstream tasks, Tune mode can save 20%-40% of memory costs without loss of accuracy.

The paper proposes a novel calculation strategy for ConvBN by changing the order of calculation. The proposed Tune mode have exact same forward/back propagation expression with the eval mode yet are faster and has lesser memory footprint, demonstrated both theoretically and empirically. This makes it a useful component in transfer learning. The authors also attributed the instability of the training using deploy mode to the scaled weight and gradient, leading to the necessity of saving part of the parameters.

This paper primarily investigates the forward computation and backpropagation process in model training or testing of three modes (i.e., Train, Eval, Deploy) of the ConvBN module (Convolution layer plus BN layer), with a focus on the instability of the Deploy mode during transfer learning. Thus, a new mode called Tune mode is proposed, which boasts the advantages of training stability and low computational time and space costs.

To verify the effectiveness of Tune mode, the authors selected 5 datasets and 12 models, and conducted extensive experiments in classification and detection tasks. The experiments show that using Tune mode for transfer learning results in a slight improvement in model performance, and a significant reduction in time and space costs. Furthermore, the authors tested the effectiveness of Tune mode in the generation of adversarial samples, thereby demonstrating the general applicability of Tune mode as a replacement for Eval mode in tasks that originally used Eval mode.

We thank all reviewers for appreciating our work. We want to highlight one contribution (the last contribution in the Introduction section) that seems to be overlooked:

Our method has been quickly integrated into open-source framework libraries like PyTorch and MMCV/MMEngine because of its evident benefit, improving the efficiency of hundreds of models for everyone using these frameworks.

For those who are using a fresh new version of PyTorch (e.g., execute conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch-nightly -c nvidia to get the latest PyTorch version), they will find our method is available in the PyTorch core library, just like those built-in modules nn.LayerNorm and nn.BatchNorm2d people use every day. Note that PyTorch acknowledges our paper in their code, but we can't include the link to the implementation due to the conference policy of double-blind review requirements.

By integrating with recent advances in deep learning compilers, our method can benefit users without changing their code. Take the following code, for example:

import torch
import torch.nn as nn

# define model architecture
class WrappedBatchNorm(nn.Module):
    def __init__(self):
        super().__init__()
        self.mod = nn.BatchNorm2d(128)
    def forward(self, x):
        return self.mod(x)

class M(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(128, 128, 1)
        self.bn1 = nn.BatchNorm2d(128)
        self.conv2 = nn.Conv2d(128, 128, 1)
        self.nested = nn.Sequential(
            nn.BatchNorm2d(128),
            nn.Conv2d(128, 128, 1),
        )
        self.wrapped = WrappedBatchNorm()

    def forward(self, x):
        # self.conv1 + self.bn1 consist of one ConvBN block
        x = self.conv1(x)
        x = self.bn1(x)
        # self.conv2 + self.nested.0 consist of one ConvBN block
        x = self.conv2(x)
        x = self.nested(x)
        # self.nested.1 + self.wrapped.mod consist of one ConvBN block
        x = self.wrapped(x)
        return x

model = M()

model.eval().cuda()

input = torch.randn(64, 128, 224, 224).cuda()

# measure the benefit of enabling Tune Mode ConvBN Block

from torch._inductor import config as inductor_config
# enable the configuration to use our method. The proposed "Tune Mode ConvBN Block" is named "efficient conv bn eval" in PyTorch.
for enabled in (True, False):
    print(f"Tune Mode ConvBN Block enabled: {enabled}")
    torch._dynamo.reset()
    torch.cuda.reset_peak_memory_stats()
    inductor_config.efficient_conv_bn_eval_fx_passes = enabled
    compiled_model = torch.compile(model)

    # warm up cuda devices by computing both forward and backward
    compiled_model(input).sum().backward()

    # measure the time and memory cost
    import time
    start_time = time.time()
    start_memory = torch.cuda.memory_allocated()
    for i in range(10):
        compiled_model(input).sum().backward()
    time_cost = time.time() - start_time
    memory_cost = torch.cuda.max_memory_allocated() - start_memory
    print(f"time cost: {time_cost} seconds")
    print(f"memory cost: {memory_cost / 1024 / 1024} MB")

The code produces the following output (using RTX 4090 GPU), which shows the benefit of our method (58% less time cost and 25% less memory cost):

Tune Mode ConvBN Block enabled: True
time cost: 1.0495944023132324 seconds
memory cost: 9432.76513671875 MB
Tune Mode ConvBN Block enabled: False
time cost: 2.4892385005950928 seconds
memory cost: 12569.3427734375 MB

The only change to use our method is a configuration flag inductor_config.efficient_conv_bn_eval_fx_passes = True.

Moreover, the above code shows that our method can directly benefit anyone using PyTorch without changing their code, even if the ConvBN patterns are very complicated (there are three ConvBN blocks in the above code: self.conv1 + self.bn1; self.conv2 + self.nested.0; self.nested.1 + self.wrapped.mod. The latter two pairs are non-trivial to discover and optimize).

To the best of our knowledge, our paper is among the first batch of papers to bridge pioneering deep learning compilers with machine learning algorithms, so that the whole PyTorch community can easily enjoy the efficiency gain brought by our method.

We hope the above clarification is clear to reviewers. Should you have any other questions, we will be happy to provide more information.

Dear Reviewers,

We are writing to provide an update on the practical applications of our method.

Recently, we have received positive feedback from the open-source community regarding our method's application in rotated object detection. Notably, it has helped reduce the memory footprint from 17.12GB to 11.57GB, achieving a 32% memory saving without affecting the model's performance. Furthermore, our method is compatible with over 80% (90 out of 111) configurations in the rotated object detection library MMRotate. This efficiency gain requires only a minimal change, specifically adding the command line option --cfg-options efficient_conv_bn_eval="[backbone]", as detailed in our supplementary material.

These new use cases reinforce the contribution we outlined in our previous comment, demonstrating how our method can seamlessly integrate with users' existing code through deep learning compilers, without the need for altering network structure definitions.

We are eager to hear your thoughts and feedback on these responses. Should you have any questions or require further information, please do not hesitate to contact us. Your insights are invaluable to the advancement of our work.

Thank you for your time and consideration.

Sincerely,

Authors of submission 3142

Four experts reviewed the paper, and all were positive about it. The rebuttal addressed the reviewers' questions and highlighted that they had incorporated the method into pytorch.

Accept (spotlight)