title: Why bottleneck structure is slower and costs more memory when training a network comparing to standard convolution?
body: <p>I'm using vnet to train a model. I want to train the model faster with less memory. So I replaced the standard 3x3 convolution to a combination of [1x1, 3x3 , 1x1] convolutions. The first 1x1 conv will reduce the channel to 1/N to reduce memory cost. The code is as follows.</p>
  
  <p>The first two classes is bottleneck structure and standard convolution. When I replace the standard convolution to bottleneck structure, <strong>although the model size and flops decrease, the real GPU memory cost and training time increase.</strong></p>
  
  <p>For example, I got :</p>
  
  <pre><code>Using standard convolution..........
  Total parameters : 10,052,609 float, model size : 39,268.00390625M
  191.78 GFLOPs
  end : 10.62517523765564s
  Max memory allocated : 3818.25341796875M
  
  Using bottleneck...........
  Total parameters : 1,145,061 float, model size : 4,472.89453125M
  16.05 GFLOPs
  end : 16.890745162963867s
  Max memory allocated : 4408.35107421875 M
  </code></pre>
  
  <p>However, in inference stage, the bottleneck structure can accelerate the network to some extend.</p>
  
  <p>Does anyone know why this happens and how to accelerate the network both in training and inference stage?</p>
  
  <p>Code : </p>
  
  <pre class="lang-py prettyprint-override"><code>import torch
  import torch.nn as nn
  import torch.nn.functional as F
  
  def groupNorm(channel, num_groups=16):
      return nn.GroupNorm(num_groups=num_groups, num_channels=channel)
  
  Norm = nn.BatchNorm3d
  
  BottleNeck_Ratio = 4
  
  
  class BottleNeck(nn.Module):
      def __init__(self, in_channels, out_channels, kernel_size, drop, stride=1, padding=1, N=BottleNeck_Ratio):
          super(BottleNeck, self).__init__()
  
          self.conv_1 = nn.Conv3d(in_channels=in_channels, out_channels=out_channels // N, kernel_size=1, stride=1)
          self.conv_2 = nn.Conv3d(in_channels=out_channels // N, out_channels=out_channels // N, kernel_size=kernel_size,
                                  stride=stride, padding=padding)
          self.conv_3 = nn.Conv3d(in_channels=out_channels // N, out_channels=out_channels, kernel_size=1, stride=1)
  
          self.norm = Norm(out_channels)
          self.relu = nn.ReLU()
          self.drop = nn.Dropout3d(drop)
  
      def forward(self, input):
          x = self.conv_1(input)
          x = self.conv_2(x)
          x = self.conv_3(x)
          return self.drop(self.relu(self.norm(x)))
  
  
  class CBR(nn.Module):
      def __init__(self, in_channels, out_channels, kernel_size, drop, stride=1, padding=1):
          super(CBR, self).__init__()
  
          self.conv = nn.Conv3d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size,
                                stride=stride, padding=padding)
          self.norm = Norm(out_channels)
          self.relu = nn.ReLU()
          self.drop = nn.Dropout3d(drop)
  
      def forward(self, input):
          return self.drop(self.relu(self.norm(self.conv(input))))
  
  
  ConvBnReluDrop = BottleNeck
  
  
  class ResidualDown(nn.Module):
      def __init__(self, in_channels, out_channels, kernel_size, drop, conv_nums, down=True):
          super(ResidualDown, self).__init__()
  
          if down:
              self.down = ConvBnReluDrop(in_channels, out_channels, kernel_size=2, stride=2, padding=0, drop=drop)
          else:
              self.down = ConvBnReluDrop(in_channels, out_channels, kernel_size=3, stride=1, padding=1, drop=drop)
  
          self.convs = nn.ModuleList()
          for i in range(conv_nums):
              self.convs.append(ConvBnReluDrop(out_channels, out_channels, kernel_size, drop))
  
          self.has_down = down
  
      def forward(self, x):
          # downsample
          res = self.down(x)
          # convolution
          out = res
          for conv in self.convs:
              out = conv(out)
          # residual
          return out + res
  
  
  class ResidualUp(nn.Module):
      def __init__(self, in_channels, out_channels, kernel_size, drop, conv_nums, up=True):
          super(ResidualUp, self).__init__()
          if up:
              self.deconv = nn.ConvTranspose3d(in_channels, out_channels, kernel_size=2, stride=2)
          else:
              self.deconv = ConvBnReluDrop(in_channels, out_channels, kernel_size=3, stride=1, padding=1, drop=drop)
  
          self.convs = nn.ModuleList()
          self.convs.append(ConvBnReluDrop(2 * out_channels, out_channels, kernel_size, drop))
          for i in range(conv_nums - 1):
              self.convs.append(ConvBnReluDrop(out_channels, out_channels, kernel_size, drop))
  
      def forward(self, big, small):
          x = self.deconv(small)
          # interpolate to prevent size not match
          x = F.interpolate(x, big.size()[-3:], mode='trilinear', align_corners=False)
          # save x as residual, [out_ch]
          res = x
  
          # skip connection, concat and conv to small's channel
          # [2*out_ch] =&gt; [out_ch]
          x = torch.cat([big, x], 1)
          for conv in self.convs:
              x = conv(x)
          return x + res
  
  
  class VBNet(nn.Module):
      def __init__(self, in_ch=1, nclass=1, drop=0.01, level=5, bn='batch', bottleneck=False):
          super(VBNet, self).__init__()
          # levels
          self.level = level
          # Normalization layer
          global Norm
          if bn == 'batch':
              Norm = nn.BatchNorm3d
          elif bn == 'group':
              Norm = groupNorm
          # elif bn == 'syncbn':
          #     Norm = SyncBN3d
          else:
              raise Exception("Error for bn")
  
          global ConvBnReluDrop
          if bottleneck:
              ConvBnReluDrop = BottleNeck
          else:
              ConvBnReluDrop = CBR
  
  
          # down 2
          self.downs = nn.ModuleList()
          self.downs.append(ResidualDown(in_ch, 16, 3, drop, 1, False))
          self.downs.append(ResidualDown(16, 32, 3, drop, 2))
  
          # down layers
          channels = 32
          for i in range(level - 2):
              self.downs.append(ResidualDown(channels, channels * 2, 3, drop, 3))
              channels *= 2
  
          # up layers
          self.ups = nn.ModuleList()
          for i in range(level - 3):
              self.ups.append(ResidualUp(channels, channels // 2, 3, drop, 3))
              channels = channels // 2
  
          # up 2
          self.ups.append(ResidualUp(channels, channels // 2, 3, drop, 2))
          channels = channels // 2
          self.ups.append(ResidualUp(channels, channels // 2, 3, drop, 1, False))
          channels = channels // 2
  
          # classifier
          self.classifier = nn.Conv3d(channels, nclass, kernel_size=1)
  
      def forward(self, x): # 4,472.89453125M
          outs = []
          for layer in self.downs:
              x = layer(x)
              outs.append(x)
  
          small = outs[-1]
          for i in range(len(self.ups)):
              layer = self.ups[i]
              big = outs[self.level - i - 2]
              small = layer(big, small)
  
          out = self.classifier(small)
          return out
  
  def get_net_size(net):
      params = list(net.parameters())
      k = 0
      for i in params:
          l = 1
          for j in i.size():
              l *= j
          k = k + l
      s = ("Total parameters : {:,} float, model size : {:,}M".format(k, k * 4 / 1024))
      return s
  
  if __name__ == '__main__':
      # count_ops is taken from : https://github.com/1adrianb/pytorch-estimate-flops/blob/master/pthflops/ops.py
      import count_ops
      import os
      import time
  
      os.environ['CUDA_VISIBLE_DEVICES'] = '1'
  
      # 4003728896
      print("Using standard convolution..........")
      a = torch.randn(6, 1, 32, 128, 128)
      net = VBNet(bn='batch', bottleneck=False)
      print(get_net_size(net))
      print(count_ops(net, a))
  
      net = net.cuda()
      start = time.time()
      for i in range(10):
          a = torch.randn(6, 1, 32, 128, 128).cuda()
          b = net(a)
          b.sum().backward()
      print('end : {}s'.format(time.time() - start))
      print("Max memory allocated : {}M".format(torch.cuda.max_memory_allocated(0) / (1024.**2)))
  
      # 4543840768 4622491136
      print("\nUsing bottleneck...........")
      # torch.cuda.reset_max_memory_allocated(0)
      a = torch.randn(6, 1, 32, 128, 128)
      net = VBNet(bn='batch', bottleneck=True)
      print(get_net_size(net))
      print(count_ops(net, a))
  
      net = net.cuda()
      start = time.time()
      for i in range(10):
          a = torch.randn(6, 1, 32, 128, 128).cuda()
          b = net(a)
          b.sum().backward()
      print('end : {}s'.format(time.time() - start))
      print("Max memory allocated : {} M".format(torch.cuda.max_memory_allocated(0) / (1024.**2)))
  
  </code></pre>