title: nn.Conv1d consumes too much GPU memory
body: <p>I have a ResBlock as below, which can change feature vector length as well. However, It consumes <strong>way too much GPU memory</strong>. In fact, one ResBlock like this alone can consume as much as <strong>2.3GB</strong> of GPU memory, which causes CUDA_OUT_OF_MEMORY all the time.</p>
  <ul>
  <li>Typical input size(batch size included): (65536, 256) or (65536, 63)</li>
  <li>Typical output size(per ResBlock): (65536, 256)</li>
  <li>The UpwardsConv1d module can change feature vector length with convolution.</li>
  </ul>
  <p>You might think the batch size is too big, but a <strong>linear layer</strong> can handle that very well, which only consumes around <strong>100MB</strong> of GPU memory per layer. There's no way that the Conv1d layer can't handle that with significantly fewer trainable parameters.</p>
  <ul>
  <li>nn.Linear(63, 256) trainable parameters: 16384</li>
  <li>ResBlock(63, 256) trainable parameters: 104</li>
  </ul>
  <pre class="lang-py prettyprint-override"><code>class UpwardsConv1d(nn.Module):
      &quot;&quot;&quot;
      Increase feature vector length by flattening dimentions.
      &quot;&quot;&quot;
  
      def __init__(self, size_in, size_out=None, size_h=8, k_size=3):
          super().__init__()
          self.size_in = size_in
          if size_out is None:
              self.size_out = size_in
          else:
              self.size_out = size_out
          self.size_h = size_h
          self.k_size = k_size
          factor = math.ceil(self.size_out/self.size_in)
  
          self.conv0 = nn.Conv1d(self.size_h, factor,
                                 self.k_size, padding=(self.k_size-1)//2)
          self.f = nn.Flatten(1, 2)
          self.conv1 = nn.Conv1d(1, 1, self.size_in*factor + 1 - self.size_out)
  
      def forward(self, x):
          x = self.conv0(x)
          x = self.f(x)
          x = x.unsqueeze(-2)
          x = self.conv1(x)
          return x
  
  
  class ResBlock(nn.Module):
      &quot;&quot;&quot;
      Act pretty much like a nn.Linear module, but uses convolution so have fewer trainable parameters.
      &quot;&quot;&quot;
  
      def __init__(self, size_in, size_out=None, size_h=2, k_size=3):
          super().__init__()
          self.size_in = size_in
          if size_out is None:
              self.size_out = size_in
          else:
              self.size_out = size_out
          self.size_h = size_h
          self.k_size = k_size
  
          self.conv0 = nn.Conv1d(1, self.size_h, self.k_size,
                                 padding=(self.k_size-1)//2)
          if self.size_in == self.size_out:
              self.conv1 = nn.Conv1d(self.size_h, 1, (self.k_size-1)//2)
          else:
              self.conv1 = UpwardsConv1d(
                  self.size_in, self.size_out, self.size_h, self.k_size)
  
      def forward(self, x):
          x = x.unsqueeze(-2)
          x = self.conv0(x)
          x = self.conv1(x)
          x = x.squeeze(-2)
          return x
  </code></pre>